一、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,打开后界面如下所示: