torch 网络模型转换onnx格式，并可视化

1. 构建lenet5 网络

import torch.nn as nn
import torch.nn.functional as F
import torch
 
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
class LeNet(nn.Module):
    def __init__(self,class_num=10,input_shape=(1,32,32)):
        super(LeNet, self).__init__()
        self.conv1 = nn.Sequential(                  #input_size=(1*28*28)
            nn.Conv2d(1, 6, 5, 1, 2),                #padding=2保证输入输出尺寸相同
            nn.ReLU(),                               #input_size=(6*28*28)
            nn.MaxPool2d(kernel_size=2, stride=2),   #output_size=(6*14*14)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(6, 16, 5),                     #padding=1输出尺寸变化
            nn.ReLU(),                               #input_size=(16*10*10)
            nn.MaxPool2d(2, 2)                       #output_size=(16*5*5)
        )

        self.fc1 = nn.Sequential(
            nn.Linear(16 * ((input_shape[1]//2-4)//2)  * ((input_shape[2]//2-4)//2), 120),
            nn.ReLU()
        )
        self.fc2 = nn.Sequential(
            nn.Linear(120, 84),
            nn.ReLU()
        )
        self.fc3 = nn.Linear(84, class_num)

                                                   # 定义前向传播过程，输入为x
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
                                                    # nn.Linear()的输入输出都是维度为一的值，所以要把多维度的tensor展平成一维
        x = x.view(x.size()[0], -1)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)

2. 转为onnx格式

input_shape = (1,100,100)   #输入数据
model = LeNet(input_shape=input_shape)

torch.save(model, './model_para.pth')
torch_model = torch.load("./model_para.pth") # pytorch模型加载
batch_size = 1  #批处理大小


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

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

3. 通过netron查看网络结构

3.1 netron安装

pip install netron

3.2 netron可视化

import netron
onnx_path = "lenet5.onnx"
netron.start(file=onnx_path, log=False, browse=True)