如何通过cv2.dnn.readNetFromTorch()读取PyTorch模型文件

问题根本原因 #

OpenCV 4.1.0版本的cv2.dnn.readNetFromTorch()接口仅支持原生Lua语言编写的Torch7框架导出的模型，不支持Python生态的PyTorch框架导出的任何格式模型。你此前调整文件后缀、切换保存state_dict还是完整模型的操作无法解决问题，本质是两种框架的序列化格式完全不兼容，报错中的「Unsupported Lua type」就是因为OpenCV误将PyTorch序列化文件当做Lua格式解析，识别失败。

可行解决方案 #

方案1：转ONNX格式后读取（最兼容4.1.0版本） #

该方案无需升级OpenCV，是适配当前环境的最优选择：

• 第一步：将PyTorch模型导出为ONNX格式，示例代码如下：

import torch
# 加载自定义模型
model = 你的模型类()
model.load_state_dict(torch.load("/home/aktaseren/people-opencv/pidxx.pt"))
model.eval()
# 构造和模型输入维度一致的伪输入，示例为单张3通道224*224的图像，可根据实际输入修改
dummy_input = torch.randn(1, 3, 224, 224)
# 导出ONNX，opset_version选11适配4.1.0版本的兼容性最好
torch.onnx.export(model, dummy_input, "/home/aktaseren/people-opencv/pidxx.onnx", opset_version=11)

• 第二步：用OpenCV直接读取ONNX模型：

import cv2
net = cv2.dnn.readNetFromONNX("/home/aktaseren/people-opencv/pidxx.onnx")

如果模型包含自定义算子，需要提前将自定义算子转为ONNX支持的标准算子组合。

方案2：升级OpenCV版本后读取TorchScript格式 #

如果不希望转ONNX，可以升级OpenCV到4.5及以上版本，该版本开始支持读取PyTorch导出的TorchScript格式模型：

• 第一步：导出TorchScript格式模型，无动态分支的模型可采用trace方式导出：

import torch
model = 你的模型类()
model.load_state_dict(torch.load("/home/aktaseren/people-opencv/pidxx.pt"))
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save("/home/aktaseren/people-opencv/pidxx_traced.pt")

• 第二步：用升级后的OpenCV读取模型：

import cv2
net = cv2.dnn.readNetFromTorch("/home/aktaseren/people-opencv/pidxx_traced.pt")

方案3：保留现有环境，拆分推理流程 #

如果必须保留OpenCV 4.1.0版本且不想修改模型导出逻辑，可放弃用OpenCV的dnn模块做推理：用PyTorch加载模型完成前处理、推理步骤，再将推理结果传入OpenCV完成后续图像操作即可，该方案仅要求运行环境同时安装PyTorch和OpenCV。

内容的提问来源于stack exchange，提问作者Eren A.