如何加载 ONNX 文件并使用它在 PyTorch 中进行 ML 预测？答案

【问题标题】：How to load an ONNX file and use it to make a ML prediction in PyTorch?如何加载 ONNX 文件并使用它在 PyTorch 中进行 ML 预测？
【发布时间】：2023-01-04 01:05:11
【问题描述】：

下面是源代码，我用它来加载.pth文件并进行多类图像分类预测。

model = Classifier()    # The Model Class.
model.load_state_dict(torch.load('<PTH-FILE-HERE>.pth'))
model = model.to(device)
model.eval()

# prediction function to test images
def predict(img_path):
    image = Image.open(img_path)
    resize = transforms.Compose(
                    [ transforms.Resize((256,256)), transforms.ToTensor()])             
    image = resize(image)
    image = image.to(device)
    y_result = model(image.unsqueeze(0))
    result_idx = y_result.argmax(dim=1)
    print(result_idx)

我使用 torch.onnx.export 将 .pth 文件转换为 ONNX 文件。

现在，如何通过单独使用 ONNX 文件而不使用 .pth 文件来编写类似于上述脚本的预测脚本。？有可能这样做吗？

【问题讨论】：

标签： pytorch onnx pth

【解决方案1】：

您可以使用 ONNX 运行时。

# !pip install onnx onnxruntime-gpu 
import onnx, onnxruntime

model_name = 'model.onnx'
onnx_model = onnx.load(model_name)
onnx.checker.check_model(onnx_model)

image = Image.open(img_path)
resize = transforms.Compose(
                [ transforms.Resize((256,256)), transforms.ToTensor()])             
image = resize(image)
image = image.unsqueeze(0) # add fake batch dimension
image = image.to(device)

EP_list = ['CUDAExecutionProvider', 'CPUExecutionProvider']

ort_session = onnxruntime.InferenceSession(model_name, providers=EP_list)

def to_numpy(tensor):
      return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

# compute ONNX Runtime output prediction
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(image)}
ort_outs = ort_session.run(None, ort_inputs)

max = float('-inf')
max_index = -1
for i in range(0, len(ort_outs[0][0])):       
   if(ort_outs[0][0][i] > max):    
       max = ort_outs[0][0][i]
       max_index = i
print(max_index)

详细解说可以关注tutorial。

通常，使用 onnx 的目的是在不同的框架中加载模型并在那里运行推理，例如PyTorch -> ONNX -> TensorRT。

【讨论】：

从 ORT 1.9 开始，需要在实例化 InferenceSession 时显式设置 providers 参数。例如， onnxruntime.InferenceSession(model_name , providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']

【解决方案2】：

如何输出置信度/比率？

【讨论】：