使用 onnx 模型的 OpenCV pytorch 预测误差

Question

我的代码在这里。我制作了一个细胞预测模型并将其转换为 onnx，然后加载 OpenCV 以使用 OpenCV 进行预测，但出了点问题

库

import cv2
import torchvision.models as models
import torch.onnx
import torchvision.transforms as transforms
import numpy as np

Onnx 和模型

original_model = models.resnet50(pretrained=True)

opencv_net = cv2.dnn.readNetFromONNX('resnet50.onnx')

Opencv 预测

opencv_net = cv2.dnn.readNetFromONNX('resnet50.onnx')
input_img=cv2.imread('image.bmp',cv2.COLOR_BGR2GRAY)
input_img=input_img.astype(np.float32)
input_img=cv2.resize(input_img,(256,256))

mean=np.array([0.485, 0.456, 0.406]) * 255.0
scale=1/255.0
std=[0.229, 0.224, 0.225]
input_blob = cv2.dnn.blobFromImage(
    image=input_img,
    scalefactor=scale,
    size=(224, 224),  # img target size
    mean=mean,
    #swapRB=True,  # BGR -> RGB
    crop=True  # center crop
)

input_blob[0] /= np.asarray(std, dtype=np.float32).reshape(3, 1, 1)
print("Input blob shape: {}\n".format(input_blob.shape))
preproc_img=input_blob

opencv_net.setInput(preproc_img)
out = opencv_net.forward()
print("OpenCV DNN prediction: \n")
print("* shape: ", out.shape)

test_class_id = np.argmax(out)

test_labels=opencv_net.getLayerNames()
#print((test_labels))

confidence = out[0][test_class_id]
print("* class ID: {}, label: {}".format(test_class_id, test_labels[test_class_id]))
print("* confidence: {:.4f}".format(confidence))

Opencv 预测输出

OpenCV DNN prediction: 

* shape:  (1, 2)
* class ID: 1, label: 323
* confidence: 8.4153

**!!!问题来了！！！ Pytorch 推理**

original_model.eval()
preproc_img = torch.FloatTensor(preproc_img)
# inference
out = original_model(preproc_img)

print("\nPyTorch model prediction: \n")
print("* shape: ", out.shape)

test_class_id = torch.argmax(out, axis=1).item()
print("* class ID: {}, label: {}".format(test_class_id, test_labels[test_class_id]))

confidence = out[0][test_class_id]
print("* confidence: {:.4f}".format(confidence.item()))

错误？

* shape:  torch.Size([1, 1000])
Traceback (most recent call last):
  File "X.py", line 121, in <module>
    print("* class ID: {}, label: {}".format(test_class_id, test_labels[test_class_id]))
IndexError: tuple index out of range

Process finished with exit code 1

当我使用相同的 img 和相同的模型时，这怎么可能？

Answer 1

我解决了这个问题。问题在于模型。我在这里使用的是 resnet 50 pretrained，但我需要我的模型，所以我用这些线来解决它并且它的工作。

model = models.resnet50(pretrained = True)
model.fc = nn.Linear(in_features=2048, out_features=2, bias=True)
weights = torch.load('model_best.pth',map_location ='cpu')
model.load_state_dict(weights)
model.eval()

给我同样的 opencv 输出

PyTorch model prediction: 

* shape:  torch.Size([1, 2])
* class ID: 1, label: 323
* confidence: 8.4153