pytorch训练的.pth模型格式转换
·
深度网络模型再工程化调用时存在模型不识别的问题,主要时推理的程序有不同的语言编写的,本文介绍使用pytorch训练的模型格式.pth转化为.pt格式的代码。
pytorch网络代码地址:https://github.com/yassouali/pytorch_segmentation ,安装请参考《pytorch框架下语义分割训练实践(一)》。
1 .pth模型转换为.pt模型
我是用fcn训练得到一个语言分割模型,checkpoint-epoch100.pth,使用inference.py文件可以正常调用,但用c++去调用的却不能正常载入。
使用下面的python脚本将.pth模型转换为.pt格式。
import torch
import torchvision
from models import fcn
model=torchvision.models.vgg16()
state_dict = torch.load("./checkpoint-epoch100.pth")
#print(state_dict)
model.load_state_dict(state_dict,False)
model.eval()
x = torch.rand(1,3,128,128)
ts = torch.jit.trace(model, x)
ts.save('fcn_vgg16.net')
注意很多人在转换的时候报错是因为:model.load_state_dict(state_dict)后面没用False参数,如下图所示。
2. .pth模型转化为.onnx模型
如需使用opencv来加载模型,则需将.pth转化为.onnx格式的模型。a.先安装onnx,使用命令:pip install onnx;b.使用以下命令转为.onnx模型
import io
import torch
import torch.onnx
import torchvision
from models import fcn
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
def test():
model=torchvision.models.vgg16()
pthfile = r'./checkpoint-epoch100.pth'
loaded_model = torch.load(pthfile, map_location='cpu')
# try:
# loaded_model.eval()
# except AttributeError as error:
# print(error)
#model.load_state_dict(loaded_model['state_dict'])
# model = model.to(device)
#data type nchw
dummy_input1 = torch.randn(1, 3, 244, 244)
# dummy_input2 = torch.randn(1, 3, 64, 64)
# dummy_input3 = torch.randn(1, 3, 64, 64)
input_names = [ "actual_input_1"]
output_names = [ "output1" ]
# torch.onnx.export(model, (dummy_input1, dummy_input2, dummy_input3), "C3AE.onnx", verbose=True, input_names=input_names, output_names=output_names)
torch.onnx.export(model, dummy_input1, "fcn.onnx", verbose=True, input_names=input_names, output_names=output_names)
if __name__ == "__main__":
test()
更多推荐
已为社区贡献6条内容
所有评论(0)