Torch network model is converted to onnx format and visualized

1. structure lenet5 The Internet

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 Make sure the input and output dimensions are the same 
            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 Output size change 
            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)

                                                   #  Define the forward propagation process , Input is x
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
                                                    # nn.Linear() The input and output of are all values with one dimension , So we need to put the multi-dimensional tensor Flatten into one dimension 
        x = x.view(x.size()[0], -1)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)

2. To onnx Format

input_shape = (1,100,100)   # input data 
model = LeNet(input_shape=input_shape)

torch.save(model, './model_para.pth')
torch_model = torch.load("./model_para.pth") # pytorch Model loading 
batch_size = 1  # Batch size 


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

x = torch.randn(batch_size,*input_shape)		    #  Generating tensor 
print (x.shape)
export_onnx_file = "lenet5.onnx"		        	#  Purpose ONNX file name 
torch.onnx.export(torch_model,
                   x,
                   export_onnx_file,
                   opset_version=10,
                   do_constant_folding=True,	#  Whether to perform constant folding optimization 
                   input_names=["input"],		#  Enter the name 
                   output_names=["output"],	#  Output name 
                   dynamic_axes={"input":{0:"batch_size"},		#  Batch variables 
                                   "output":{0:"batch_size"}})

3. adopt netron Look at the network structure

3.1 netron install

pip install netron

3.2 netron visualization

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