import torch
import torchvision
import cv2
import onnx
import numpy as np
import matplotlib.pyplot as plt
print(torch.__version__)
print(cv2.__version__)
print(np.__version__)

def init_model(model_name):
    if model_name=='alexnet':
        model = torchvision.models.alexnet(pretrained=True)
    if model_name=='densnet':
        model = torchvision.models.densenet121(pretrained=True)
    if model_name=='resnet':
        model = torchvision.models.resnet50(pretrained=True)
    if model_name=='mobilenet':    
        model = torchvision.models.mobilenet_v2(pretrained=True)
    if model_name=='squeezenet':
        model = torchvision.models.squeezenet1_1(pretrained=True)
    if model_name=='inception':
        model = torchvision.models.inception_v3(pretrained=True)
    if model_name=='googlenet':
        model = torchvision.models.googlenet(pretrained=True)
    if model_name=='vgg16':
        model = torchvision.models.vgg16(pretrained=True)
    if model_name=='vgg19':
        model = torchvision.models.vgg19(pretrained=True)
    if model_name=='shufflenet':
        model = torchvision.models.shufflenet_v2_x1_5(pretrained=True)
    model.eval()
    if model_name=='inception':
        dummy = torch.randn(1,3,299,299)
    else:
        dummy = torch.randn(1,3,224,224)
    return model, dummy


model, dummy = init_model('inception')

onnx_name = 'exported.onnx'
torch.onnx.export(model, dummy, onnx_name)

# 载入onnx模块
model_ = onnx.load(onnx_name)
#检查IR是否良好
onnx.checker.check_model(model_)

# opencv dnn加载
net = cv2.dnn.readNetFromONNX(onnx_name)

img = r"C:\Users\admin\Pictures\cat.jpg"
frame = cv2.imread(img)
classes = None
class_file = r"F:\opencv\sources\samples\data\dnn\classification_classes_ILSVRC2012.txt"
with open(class_file, 'rt') as f:
    classes = f.read().rstrip('\n').split('\n')
    
# Create a 4D blob from a frame.
inpWidth = dummy.shape[-2]
inpHeight = dummy.shape[-2]
blob = cv2.dnn.blobFromImage(frame, 
                             size=(inpWidth, inpHeight), crop=False)

# Run a model
net.setInput(blob)
out = net.forward()

# Get a class with a highest score.
out = out.flatten()
classId = np.argmax(out)
confidence = out[classId]

# Put efficiency information.
t, _ = net.getPerfProfile()
label = 'Inference time: %.2f ms' % (t * 1000.0 / cv2.getTickFrequency())
print(label)
cv2.putText(frame, label, (0, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0))

# Print predicted class.
label = '%s: %.4f' % (classes[classId] if classes else 'Class #%d' % classId, confidence)
print(label)
cv2.putText(frame, label, (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0))
winName = 'onnx'

cv2.imshow(winName, frame)
cv2.waitKey(0)
cv2.destroyAllWindows()

1.5.1+cpu
4.4.0
1.16.6
Traceback (most recent call last):

  File "E:\OneDrive - Dezhkeda\Files\MachineLearning\python\pytorch\onnx_test.py", line 71, in <module>
    net = cv2.dnn.readNetFromONNX(onnx_name)

error: OpenCV(4.4.0) F:\opencv-4.4.0\modules\dnn\src\onnx\onnx_importer.cpp:262: error: (-204:Requested object was not found) Blob x.1 not found in const blobs in function 'cv::dnn::dnn4_v20200609::ONNXImporter::getBlob'