pth转onnx

前段时间做项目需要将pth转成onnx模型，这里将经验分享出来。

首先你得看你的pth模型是只有参数还是说既有参数，也有模型结构，怎么看呢？将你的pth模型在netron中打开Netron，如果模型是黑白的，就说明只有参数，没有模型结构，如图一。如果是彩色的，就是既有参数也有模型结构的，如图二。

下面直接上代码（只有参数的pth）：

import torch
import nets.facenet as models
from utils.config import cfg_mnet

torch_model = torch.load("facenet_mobilenet.pth", map_location= 'cpu') # pytorch模型加载
model = models.Facenet(mode='val')
#model.load_state_dict(torch_model)

batch_size = 1  #批处理大小
input_shape = (3, 160, 160)   #输入数据,改成自己的输入shape

# #set the model to inference mode
model.eval()

x = torch.randn(batch_size, *input_shape)  # 生成张量
export_onnx_file = "facenet_mobilenet.onnx"          # 目的ONNX文件名
torch.onnx.export(model,
                    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"}})

整体思路就是：因为我们模型只有参数，所以要加上模型架构。

对于既有参数有模型架构的，网络上有很多，这里就不说了。

图一：

图二：