pytorch pth模型转onnx

转换方法：

使用pytorch自带转换函数转换

import torch

# 实例你的模型
net = MyModel()
print('load static dict')
# 加载参数！ （如果忘记的话，，，就重来呗）
net.load_state_dict(torch.load("net.pth", map_location=torch.device('cpu')))
net.eval().cuda()
print('to cuda')

input1 = torch.randn(1, 3, 112,112).cuda()
input_names = [ "input"]
output_names = [ "output" ]
   
torch.onnx.export(net, input1, "out.onnx", verbose=True, input_names=input_names, output_names=output_names)

# 如果有多个输入可以用元组的形式包起来  如：
input2 = torch.randn(1,3,112,112).cuda()
torch.onnx.export(net, (input1,input2), "out.onnx", verbose=True, input_names=input_names, output_names=output_names)

验证

验证下onnx模型是否可用

import onnxruntime as ort
import numpy as np
import time

x = np.random.rand(1,3,112,112).astype(np.float32)
ort_sission = ort.InferenceSession('out.onnx')
onnx_output = ort_sission.run(None,{ort_sission.get_inputs()[0].name:x})
print(onnx_output)
# 这里可以对比使用onnx模型的输出 和pth的输出

# 统计时间
begin = time.time()
print('begin')
for  i in range(100):
    onnx_output = ort_sission.run(None,{ort_sission.get_inputs()[0].name:x})
end = time.time()
print('emd')
print(1/(end-begin)*100)
# 使用sleep，在sleep时运行时开辟的内存并不会释放掉，所以可以更方便的观察内存使用情况
time.sleep(30)