PyTorch转ONNX模型无法在OpenCV DNN加载的问题排查

问题背景 #

OpenCV DNN不支持直接加载PyTorch的.pt格式模型，因此按照PyTorch官方流程将模型转换为ONNX格式后，ONNX Runtime可正常运行且输出与原模型一致，但使用OpenCV DNN加载时出现Pad节点相关错误。当前环境：ONNX 1.13、ONNXruntime 1.14、OpenCV 4.7。

操作步骤 #

步骤1：定义PyTorch神经网络 #

import torch
import torch.nn as nn
import torch.nn.functional as F

class ModelTeste(nn.Module):
    def __init__(self):
        super(ModelTeste, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size = 20, stride=3)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(32, 42, kernel_size=12, stride=2,)
        self.relu2 = nn.ReLU()
        self.pool = nn.MaxPool2d(kernel_size=2)
        self.conv3 = nn.Conv2d(42, 24, kernel_size=9, stride=2)
        self.relu3 = nn.ReLU()
        self.conv4 = nn.Conv2d(24, 16, kernel_size=5, stride=2)
        self.relu4 = nn.ReLU()
        self.flatten = nn.Flatten(1, 3)
        self.fc1 = nn.Linear(144, 350)
        self.relu5 = nn.ReLU()
        self.dropout = nn.Dropout(0.05)
        self.fc2 = nn.Linear(350, 350)
        self.relu6 = nn.ReLU()
        self.fc3 = nn.Linear(350, 350)
        self.relu7 = nn.ReLU()
        self.fc4 = nn.Linear(350, 400)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.conv1(F.pad(x, [8, 9, 8, 9]))
        x = self.relu1(x)
        x = self.conv2(F.pad(x, [5, 5, 5, 5]))
        x = self.relu2(x)
        x = self.pool(x)
        x = self.conv3(F.pad(x, [3, 4, 3, 4]))
        x = self.relu3(x)
        x = self.conv4(F.pad(x, [2, 2, 2, 2]))
        x = self.relu4(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu5(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.relu6(x)
        x = self.fc3(x)
        x = self.relu7(x)
        x = self.fc4(x)
        x = self.sigmoid(x)
        return x
    
modelteste = ModelTeste()
criterionteste = torch.nn.BCELoss()
optimizerteste = torch.optim.Adam(modelteste.parameters(), lr=0.001) 

步骤2：加载预训练模型状态字典 #

PATH = 'C:\\neokoros\\pytorchconversao\\statedict_modelnovo.pt'
batch_size = 64

# 加载预训练权重
modelteste.load_state_dict(torch.load(PATH))
# 设置为推理模式
modelteste.eval()

步骤3：导出为ONNX模型 #

# 创建虚拟输入
x = torch.rand(1, 1, 120, 120, requires_grad=True)

torch.onnx.export(modelteste,               # 待导出模型
                  x,                         # 模型输入
                  "model1.onnx",             # 导出路径
                  verbose=True,
                  export_params=True,        # 保存模型参数
                  opset_version=12,          # ONNX版本
                  do_constant_folding=True,  # 启用常量折叠优化
                  input_names = ['input'],   # 输入节点名称
                  output_names = ['output'], # 输出节点名称
                  dynamic_axes={'input' : {0 : 'batch_size'},    # 动态batch轴
                                'output' : {0 : 'batch_size'}})

步骤4：ONNX Runtime验证模型 #

import onnx
import onnxruntime as rt
import numpy as np

model1 = onnx.load("model1.onnx")
print(onnx.helper.printable_graph(model1.graph)) # 查看模型结构

x_out = x.detach().numpy() # 将张量转为numpy数组

sess = rt.InferenceSession("model1.onnx")
input_names = sess.get_inputs()[0].name
print(input_names)
output_names = sess.get_outputs()[0].name
print(output_names)
scores = sess.run([output_names], {input_names: x_out.astype(np.float32)})[0]
print(scores)

此步骤验证模型运行正常，输出与PyTorch原模型一致。

错误信息 #

执行net = cv2.dnn.readNet("model1.onnx")时，出现以下错误：

error: OpenCV(4.7.0) D:\a\opencv-python\opencv-python\opencv\modules\dnn\src\onnx\onnx_importer.cpp:1073: error: (-2:Unspecified error) in function 'cv::dnn::dnn4_v20221220::ONNXImporter::handleNode'

Node [Pad@ai.onnx]:(onnx_node!Pad_16) parse error: OpenCV(4.7.0) D:\a\opencv-python\opencv-python\opencv\modules\dnn\src\onnx\onnx_importer.cpp:591: error: (-5:Bad argument) Blob  not found in const blobs in function 'cv::dnn::dnn4_v20221220::ONNXImporter::getBlob'

解决思路 #

1. 替换手动Pad为卷积层内置Padding #

问题根源在于PyTorch导出的F.pad生成的独立Pad节点与OpenCV DNN兼容不佳。可以将手动补零操作整合到卷积层的padding参数中，避免生成独立Pad节点：

• 原代码：x = self.conv1(F.pad(x, [8, 9, 8, 9]))
• 修改为在定义conv1时设置padding：self.conv1 = nn.Conv2d(1, 32, kernel_size=20, stride=3, padding=((8,9), (8,9)))（对应上8、下9、左8、右9的补零）
• 同理修改其他卷积层的padding参数，移除所有F.pad调用后重新导出ONNX模型。

2. 禁用动态轴导出 #

OpenCV DNN对ONNX动态轴支持有限，即使仅batch_size动态也可能引发兼容问题。导出ONNX时删除dynamic_axes参数，导出固定batch_size的模型：

torch.onnx.export(modelteste,
                  x,
                  "model1_static.onnx",
                  verbose=True,
                  export_params=True,
                  opset_version=12,
                  do_constant_folding=True,
                  input_names = ['input'],
                  output_names = ['output'])

3. 升级OpenCV版本 #

OpenCV 4.7对ONNX Pad算子的支持存在缺陷，升级至4.8及以上版本可修复部分兼容问题，新版本DNN模块对ONNX算子的支持更完善。

4. 使用ONNX Simplifier优化模型 #

通过onnx-simplifier工具简化模型，可能将Pad节点与后续卷积节点融合，转换为OpenCV可识别的格式：

• 安装工具：pip install onnx-simplifier
• 简化模型：python -m onnxsim model1.onnx model1_simplified.onnx
• 尝试加载简化后的模型。

内容的提问来源于stack exchange，提问作者mentejoao