将pytorch模型转换为ONNX

Question

如何将火把模型转换为ONNX？我试图在Python3.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 7x型号的git回购：

​

最终的用例是在Intel的Open工具包上使用这个模型，它需要将PyTorch模型转换为ONYX。

Answer 1

在加载带泡菜的模型时，源树必须与保存模型时使用的树匹配。所以

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 (用内特恩可视化)：

​

通常，最好将权重保存为state_dict，并保留可以重构torch.nn.Module的源代码，这样您就可以安全地使用：

model.load_state_dict(torch.load('weights.pt'))

Answer 2

至于它的价值，我只能让它工作曾傑瑞操纵一些代码从另一个回购，加上答案，从@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一起工作，这是另一个火把用户回购的一个文件。我不熟悉火把，不知道这是不是一个普通的火把文件.

编辑

查看open : Yolo7上的REPO：巨蟒