OpenCV pytorch 预测错误与 onnx model

Question

我的代码在这里。 我制作了一个 model 用于预测单元格并将其转换为 onnx 然后加载 OpenCV 以预测 OpenCV 但出了点问题

图书馆

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

Onnx 和 model

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 预测 output

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 和相同的 model 时，这怎么可能？

Answer 1

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

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