Pytorch与Onnx的转换与推理
Open Neural Network Exchange（ONNX，开放神经网络交换）格式，是一个用于表示深度学习模型的标准，可使模型在不同框架之间进行转移。
一、pytorch模型保存/加载
有两种方式可用于保存/加载pytorch模型 1）文件中保存模型结构和权重参数 2）文件只保留模型权重.
1、文件中保存模型结构和权重参数
模型保存与调用方式一（只保存权重）：
保存：
torch.save(model.state_dict(), mymodel.pth)#只保存模型权重参数，不保存模型结构调用：
model = My_model(*args, **kwargs) #这里需要重新创建模型，My_model
model.load_state_dict(torch.load(mymodel.pth))#这里根据模型结构，导入存储的模型参数
model.eval()模型保存与调用方式二（保存完整模型）：
保存：
torch.save(model, mymodel.pth)#保存整个model的状态调用：
model=torch.load(mymodel.pth)#这里已经不需要重构模型结构了，直接load就可以
model.eval().pt表示pytorch的模型，.onnx表示onnx的模型,后缀名为.pt, .pth, .pkl的pytorch模型文件之间其实没有任何区别
二、pytorch模型转ONNX模型
1、文件中保存模型结构和权重参数
import torch
torch_model = torch.load("/home/pytorch/save.pth") # pytorch模型加载
#set the model to inference mode
torch_model.eval()
x = torch.randn(1,3,320,640) # 生成张量（模型输入格式）
export_onnx_file = "/home/pytorch/test.onnx" # 目的ONNX文件名
// 导出export:pt->onnx
torch.onnx.export(torch_model, # pytorch模型
x, # 生成张量（模型输入格式）
export_onnx_file, # 目的ONNX文件名
do_constant_folding=True, # 是否执行常量折叠优化
input_names=["input"], # 输入名（可略）
output_names=["output"], # 输出名（可略）
dynamic_axes={"input":{0:"batch_size"}, # 批处理变量（可略）
"output":{0:"batch_size"}}) 注：dynamic_axes字段用于批处理.若不想支持批处理或固定批处理大小,移除dynamic_axes字段即可.
2、文件中只保留模型权重
import torch
torch_model = selfmodel() # 由研究员提供python.py文件
#set the model to inference mode
torch_model.eval()
x = torch.randn(1,3,320,640) # 生成张量（模型输入格式）
export_onnx_file = "/home/pytorch/test.onnx" # 目的ONNX文件名
// 导出export:pt->onnx
torch.onnx.export(torch_model, # pytorch模型
x, # 生成张量（模型输入格式）
export_onnx_file, # 目的ONNX文件名
do_constant_folding=True, # 是否执行常量折叠优化
input_names=["input"], # 输入名（可略）
output_names=["output"], # 输出名（可略）
dynamic_axes={"input":{0:"batch_size"}, # 批处理变量（可略）
"output":{0:"batch_size"}}) 3、onnx文件操作
3.1 安装onnx,onnxruntime:
pip install onnx
pip install onnxruntime(只能用cpu)
pip install onnxruntime-gpu(gpu和cpu都能用)首先要强调的是，有两个版本的onnxruntime，一个叫onnxruntime，只能使用cpu推理，另一个叫onnxruntime-gpu，既可以使用gpu，也可以使用cpu。
如果自己安装的是onnxruntime，需要卸载后安装gpu版本。
确认一下是否可以使用gpu
注意：
```python
print(onnxruntime.get_device())
```
上面的代码给出的输出是'GPU'时，并不代表就成功了。
而要用下面的代码来验证：
```python
ort_session = onnxruntime.InferenceSession("path/model/model_name.onnx",
providers=['CUDAExecutionProvider'])
print(ort_session.get_providers())
```
当输出是:['CUDAExecutionProvider', 'CPUExecutionProvider']才表示成功了。
版本查询：NVIDIA - CUDA | onnxruntime
安装固定版本的onnxruntime:
pip install onnxruntime-gpu==1.9.0
卸载pip uninstall
3.2 加载onnx文件
# "加载load"
model=onnx.load('net.onnx')检查模型格式是否完整及正确
onnx.checker.check_model(model)3.3 打印onnx模型文件信息
session=onnxruntime.InferenceSession('net.onnx')
inp=session.get_inputs()[0]
#conv1=session.get_inputs()['conv1']
#out1=session.get_outputs()[1]
out=session.get_provider_options()
#print(inp,conv1,out1)
print(inp)
#print(out)
"打印图信息:字符串信息"
graph=onnx.helper.printable_graph(model.graph)
print(type(graph))3.4 获取onnx模型输入输出层
input=model.graph.input
output = model.graph.output
"""输入输出层"""
print(input,output)3.5 推理过程
import onnx
import onnxruntime
import torch
inputs=torch.randn(1,3,640,320)
#上述inputs仅用于测试使用，用于图片推理，应该换成自己的图片，如：
#img_path='1.jpg'#图片尺寸与onnx模型的处理尺寸保持一致
#img=cv2.imread(img_path)
#inputs=preprocess_imgae(img)#标准化等预处理操作，与源项目代码保持一致即可
#print('inputs.size():',inputs.size())
model=onnx.load('/home/pytorch_DL/test_320_640.onnx')
onnx.checker.check_model(model)
session =onnxruntime.InferenceSession('/home/pytorch_DL/test_320_640.onnx',
providers['CUDAExecutionProvider','CPUExecutionProvider'])
print('session.get_providers():',session.get_providers())
input_name = session.get_inputs()
output_name=session.get_outputs()[0].name
res=session.run([output_name],{input_name[0].name:inputs.numpy()})
参考：Pytorch与Onnx模型的保存、转换与操作_onnx转pytorch_Yuezero_的博客-CSDN博客
pytorch 模型的保存与加载方法以及使用onnx模型部署推理 | 码农家园
onnxruntime使用gpu推理 - 知乎