将pytorch模型转换为ONNX

fjnneemd 于 2022-11-09 发布在其他

关注(0)|答案(2)|浏览(220)

如何将pytorch模型转换为ONNX？我尝试在Python 3.7上使用这个方法：

import torch

model = torch.load("./yolov7x.pt")

# torch.onnx.export(model, "yolo_v7x.onnx")

即使在3行代码中的最后一行带有注解，加载时也会出现以下错误：

Traceback (most recent call last):
  File "C:\Users\convert_onx.py", line 5, in <module>
    model = torch.load("./yolov7x.pt")
  File "C:\Users\Python37\lib\site-packages\torch\serialization.py", line 594, in load
    return _load(opened_zipfile, map_location, pickle_module,**pickle_load_args)
  File "C:\Users\Python37\lib\site-packages\torch\serialization.py", line 853, in _load
    result = unpickler.load()
ModuleNotFoundError: No module named 'models'

这是我正在使用的Yolo Model 7x的git存储库：

最终使用案例是在英特尔的Open VINO工具包上使用此模型，该工具包要求将PyTorch模型转换为ONYX。

pytorch

来源：https://stackoverflow.com/questions/74041752/convert-pytorch-model-to-onnx

2条答案

按热度按时间

carvr3hs1#

当您加载pickle模型时，源代码树必须与保存模型时使用的源代码树相匹配。

ModuleNotFoundError: No module named 'models'

期望此目录位于python路径中：https://github.com/WongKinYiu/yolov7/tree/main/models
要导出到ONNX：
1.克隆存储库https://github.com/WongKinYiu/yolov7

git clone https://github.com/WongKinYiu/yolov7

1.请设置正确的路径。

import sys
sys.path.insert(0, './yolov7')

或者可以设置PYTHONPATH环境变量
1.另外，你可能需要有特定的torch版本。我已经检查过了，它似乎torch==1.8.0工作正常
示例：

import torch
import sys

sys.path.insert(0, './yolov7')

device = torch.device('cpu')
model = torch.load('yolov7x.pt', map_location=device)['model'].float()
torch.onnx.export(model, torch.zeros((1, 3, 640, 640)), 'yolov7.onnx', opset_version=12)

之后，将模型导出到ONNX（使用netron可视化）：

通常，最好将权重保存为state_dict，并保留可以重构torch.nn.Module的源代码，这样您就可以安全地用途：
model.load_state_dict(torch.load('weights.pt'))

赞(0）回复(0）举报 2022-11-09

jogvjijk2#

为了它的价值，我只能得到这个工作杰里操纵一些代码从另一个回购，加上答案从@u1234x1234

import torch
from models.experimental import attempt_load
import os

weights ='yolov7x.pt'
device = torch.device('cpu')

# Load model

model = attempt_load(weights, map_location=device) 

torch.onnx.export(model, torch.zeros((1, 3, 640, 640)), 'yolov7.onnx', opset_version=12)

似乎与此函数attempt_load一起工作，一个来自this other pytorch user repo的文件。我不熟悉pytorch，不知道这是否只是一个普通的pytorch文件。

编辑

请参阅打开VINO上Yolo 7的此REPO：https://github.com/OpenVINO-dev-contest/YOLOv7_OpenVINO_cpp-python

赞(0）回复(0）举报 2022-11-09

我来回答

将pytorch模型转换为ONNX

2条答案

相关问题

热门标签

最新问答