pth文件转换onnx生成可视化网络图_pth转onnx 可视化

参考文章：pth转onnx_小杨树718的博客-CSDN博客

软件:netron

pth:自己生成的权重文件。

实际上也可以将pth文件放到netron当中，但是得出来的图之间没有连接。示意图如下：

建一个文件叫onnx.py,然后把代码放到里面。

代码如下：

import torch.onnx
from torchvision import models
 
torch_model = torch.load("F:\pyskl-main\work_dirs\posec3d\c3d_light_gym\joint\epoch_1.pth", map_location='cpu')  # pytorch模型加载
model = models.resnet50()
# 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"}})

运行结果如下：

生成了onnx文件之后，打开网站netron，点击open model，选择onnx文件。就可以输出模型了。