pytorch转ONNX文件

加载模型

加载模型就是正常的用pytorch去读取自己的model就可以啦，我用的是YOLO模型，所以就用了下官网的读取模型的方法。

# 加载模型，这里可以用自己加载模型的方式
model = attempt_load(r"C:\Users\Burt\Desktop\test\best.pt")

# eval 遍历的时候不会改变参数
model.eval()

# GPU训练的
model = model.cuda()

定义输入

这个输入就是用来遍历模型的，数值不重要，主要是模型输入大小要对的上，我的模型输入是（1，3，640，640）的，就直接定义一个这么大的随机量就行了。

# 定义遍历模型的输入
dummy_input = torch.randint(0, 10, (1, 3, 640, 640))

定义输出模型路径

就是转换完成后模型保存的地址，像我就是直接定义"C:/Users/Burt/Desktop/test/best.onnx"

# 定义输出模型的路径
f = "path/to/model.onnx"

转换模型

# 转换模型
torch.onnx.export(model,
                  dummy_input,
                  f)

export参数：
第一个参数：model，自己用pytorch读取的模型
第二个参数：args，输入
第三个参数：f，输出的路径
还有一些其他参数可以，可以参考

https://www.yht7.com/news/199972

完整代码

import torch
import onnx

# 加载模型
model = attempt_load(r"C:\Users\Burt\Desktop\test\best.pt")  # load FP32 model
model.eval()
model = model.cuda()

#定义输入
dummy_input = torch.randint(0, 10, (1, 3, 640, 640))

# 输出
torch.onnx.export(model,
                  dummy_input,
                  "C:/Users/Burt/Desktop/test/best.onnx")

文件图