OpenCV:从onnx文件加载Net时出错



我正试图加载cv.dnn.readNetFromONNX预训练的火炬模型(准确地说是U2Net)保存为onnx

但是我收到错误:

error: OpenCV(4.1.2) /io/opencv/modules/dnn/include/opencv2/dnn/dnn.inl.hpp:349: 
error (-204:Requested object was not found) Required argument "starts" not found 
into dictionary in function 'get'

这是用Google Colab重现错误的代码:

### get U2Net implementation ###
%cd /content
!git clone https://github.com/shreyas-bk/U-2-Net
### download pre-trained model ###
!gdown --id 1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ -O /content/U-2-Net/u2net.pth
###
%cd /content/U-2-Net
### imports ###
from google.colab import files
from model import U2NET
import torch
import os
### create U2Net model from state
model_dir = '/content/U-2-Net/u2net.pth'
net = U2NET(3, 1)
net.load_state_dict(torch.load(model_dir, map_location='cpu'))
net.eval()
### pass to it a dummy input and save to onnx ###
img = torch.randn(1, 3, 320, 320, requires_grad=False)
img = img.to(torch.device('cpu'))
output_dir = os.path.join('/content/u2net.onnx')
torch.onnx.export(net, img, output_dir, opset_version=11, verbose=True)
### load the model in OpenCV ###
import cv2 as cv
net = cv.dnn.readNetFromONNX('/content/u2net.onnx')

[OpenCV =>4.1.2,平台=>Google Colab, Torch =>1.11.0 + cu113]

正如@berak所建议的,这个问题与OpenCV版本(4.1.2)有关。更新到4.5.5已经解决了这个问题。

最新更新