1.掌握pytorch模型转换到onnx模型
2.顺利运行onnx模型
3.比对onnx模型和pytorch模型的输出结果
1 . pytorch 转 onnx
pytorch 转 onnx 只需要一个函数 torch.onnx.export
torch.onnx.export(model, args, path, export_params, verbose, input_names, output_names, do_constant_folding, dynamic_axes, opset_version)
参数说明：
- model——需要导出的pytorch模型
- args——模型的输入参数，满足输入层的shape正确即可。
- path——输出的onnx模型的位置。例如’yolov5.onnx’。
- export_params——输出模型是否可训练。default=True，表示导出trained model，否则untrained。
- verbose——是否打印模型转换信息。default=False。
- input_names——输入节点名称。default=None。
- output_names——输出节点名称。default=None。
- do_constant_folding——是否使用常量折叠（不了解），默认即可。default=True。
- dynamic_axes——模型的输入输出有时是可变的，如Rnn，或者输出图像的batch可变，可通过该参数设置。如输入层的shape为（b，3，h，w），batch，height，width是可变的，但是chancel是固定三通道。
格式如下 ：
1)仅list(int) dynamic_axes={‘input’:[0,2,3],’output’:[0,1]}
2)仅dict - opset_version——opset的版本，低版本不支持upsample等操作。
import torch
import torch.nn
import onnx
model = torch.load('best.pt')
model.eval()
input_names = ['input']
output_names = ['output']
x = torch.randn(1,3,32,32,requires_grad=True)
torch.onnx.export(model, x, 'best.onnx', input_names=input_names, output_names=output_names, verbose='True')
2 . 运行onnx模型
检查onnx模型，并使用onnxruntime运行。
import onnx
import onnxruntime as ort
model = onnx.load('best.onnx')
onnx.checker.check_model(model)
session = ort.InferenceSession('best.onnx')
x=np.random.randn(1,3,32,32).astype(np.float32)
outputs = session.run(None,input = { 'input' : x })
参数说明：
- output_names: default=None
用来指定输出哪些，以及顺序
若为None，则按序输出所有的output，即返回[output_0,output_1]
若为[‘output_1′,’output_0’]，则返回[output_1,output_0]
若为[‘output_0’]，则仅返回[output_0:tensor] - input:dict
可以通过session.get_inputs().name获得名称
其中key值要求与torch.onnx.export中设定的一致。
3.onnx模型输出与pytorch模型比对
import numpy as np
np.testing.assert_allclose(torch_result[0].detach().numpu(),onnx_result,rtol=0.0001)
Original: https://blog.csdn.net/weixin_38989668/article/details/123840882
Author: 挣扎的笨鸟
Title: pytorch模型转onnx模型
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/670676/
转载文章受原作者版权保护。转载请注明原作者出处！