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()