深度学习|pth转onnx pytorch|人工智能

前段时间做项目需要将pth转成onnx模型，这里将经验分享出来。
首先你得看你的pth模型是只有参数还是说既有参数，也有模型结构，怎么看呢？将你的pth模型在netron中打开Netron，如果模型是黑白的，就说明只有参数，没有模型结构，如图一。如果是彩色的，就是既有参数也有模型结构的，如图二。
下面直接上代码（只有参数的pth）：

import torch import nets.facenet as models from utils.config import cfg_mnettorch_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"}})

整体思路就是：因为我们模型只有参数，所以要加上模型架构。
对于既有参数有模型架构的，网络上有很多，这里就不说了。
图一：