模型文件转换（pkl转pt转onnx转ncnn）

一、pkl转pt 

本人只试了pkl转pt，pth等文件转pt应该是一样的。

import torch
# An instance of your model.
from models.MIMOUNet import MIMOUNetPlus
import torchvision.transforms as transforms

model = MIMOUNetPlus()
state_dict = torch.load('model.pkl') # pkl模型文件路径
model.load_state_dict(state_dict) 

# model.cuda()
model.eval()

example = torch.rand(1, 3, 256, 256) # 与训练模型时输入大小相同

# example = example.cuda()

traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("model.pt") # 更换为自己的模型输出路径

可能需要注意的地方：

1、model.load_state_dict(state_dict) 有可能需要指定state_dict['model']属性传入，具体看pkl文件的数据保存形式。

2、model.cuda()和example = example.cuda()两行代码如果最终运行环境支持GPU调用的话就加上，不加就仅使用CPU处理。

3、某些网络最后输出的是一个list（可能包含多个4维张量），需要在模型的.py文件里指定输出最终的那一个4维张量。比如在forward里return output[2]，返回[1,3,size,size]。

 二、pt转onnx

import torch

# 转onnx
device = torch.device('cpu')
model = torch.jit.load('model.pt', map_location=device) # pt文件路径
model.eval()
# 下面根据不同模型的输入来决定
batch_size=1
channel=3
w_size=256
h_size=256
x=torch.randn(batch_size,channel,h_size,w_size,requires_grad=True).to(device)
output=model(x)
torch.onnx.export(model, x, 'model.onnx', input_names=['input'], output_names=["output"], opset_version=11)    # onnx输出路径

修改输入输出路径和size即可。

三、onnx转onnx-simplifier

直接onnx模式在进行后续的部署时容易出现问题，转换成simple模式。

安装onnx-simplifier：

pip install onnx-simplifier

输入命令转换模型：

python -m onnxsim model.onnx model_sim.onnx

转换完成后可以得到sim.onnx文件。

四、onnx-simplifier转ncnn

线上转换地址为：  https://convertmodel.com/#outputFormat=ncnn，打开后界面如下所示：