pytorch的.pkl文件转为onnx文件，亲测可行

pytorch的.pkl文件转为onnx文件，亲测可行

1、训练文件中需要有保存.pkl的代码

        torch.save(model,'./torch-save.pkl')

直接保存模型参数和模型的结构（建议这么使用）

2、torch2onnx.py

import torch


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")            #如果训练时用的GPU，必须还得使用GPU
torch_model = torch.load("torch-save.pkl") 											#加载.pkl文件
batch_size = 1  										#batch_size需要定下来，可不为1
input_shape = (3,208,976)                     												#模型的输入，根据训练时数据集的输入

# set the model to inference mode
torch_model.eval()            #切换到推理模式

x = torch.randn(batch_size,*input_shape)		
x = x.to(device)
export_onnx_file = "torch-save.onnx"					
torch.onnx.export(torch_model.module,
                    x,
                    export_onnx_file,
                    opset_version=10,
                    do_constant_folding=True,	
                    input_names=["input"],	
                    output_names=["output"],	
                    dynamic_axes={"input":{0:"batch_size"},	
                                    "output":{0:"batch_size"}})

3、亲测可行