深度学习模型显示工具netron

深度学习模型显示工具netron

APP网址

https://netron.app/

netron是很不的深度学习模型显示工具

netron支持显示大多数的深度学习模型，不支持pytorch生成的pt或者pth文件，但是将这两种文件转为onnx格式，netron是支持的

安装

其他的系统个人还没有尝试过，但是github地址中有，个人暂时只在ubuntu下安装了，很简单

 pip install netron

代码

 import torch
 import torch.nn as nn
 import netron
 ​
 # 定义一个简单的二分类网络
 class SimpleNet(nn.Module):
     def __init__(self):
         super(SimpleNet, self).__init__()
         self.conv1 = nn.Sequential(
             nn.Conv2d(in_channels=3, out_channels=50, kernel_size=3, stride=1, padding=1),
             nn.ReLU(),
             nn.MaxPool2d(2)
         )
         self.conv2 = nn.Sequential(
             nn.Conv2d(in_channels=50, out_channels=200, kernel_size=3, stride=1, padding=1),
             nn.ReLU(),
             nn.MaxPool2d(2)
         )
         self.conv3 = nn.Sequential(
             nn.Conv2d(in_channels=200, out_channels=500, kernel_size=3, stride=1, padding=1),
             nn.ReLU(),
             nn.MaxPool2d(2)
         )
 ​
         self.conv4 = nn.Sequential(
             nn.Conv2d(in_channels=500, out_channels=200, kernel_size=3, stride=1, padding=1),
             nn.ReLU(),
             nn.MaxPool2d(2)
         )
         self.conv5 = nn.Sequential(
             nn.Conv2d(in_channels=200, out_channels=50, kernel_size=3, stride=1, padding=1),
             nn.ReLU(),
             nn.MaxPool2d(2)
         )
         self.fc = nn.Sequential(
             nn.Linear(50 * 20 * 11, 5000)
         )
         self.fc1 = nn.Sequential(
             nn.Linear(5000, 2000)
         )
         self.fc2 = nn.Sequential(
             nn.Linear(2000, 50)
         )
         self.classifier = nn.Sequential(
             nn.Linear(50, 2)
         )
 ​
     def forward(self, x):
         x = torch.tensor(x, dtype=torch.float32)
         x = self.conv1(x)
         x = self.conv2(x)
         x = self.conv3(x)
         x = self.conv4(x)
         x = self.conv5(x)
         x = torch.flatten(x, start_dim=1)
         x = self.fc(x)
         x = self.fc1(x)
         x = self.fc2(x)
         x = self.classifier(x)
         return x
 ​
 d = torch.rand(1, 3, 640, 360)
 m = SimpleNet()
 o = m(d)
 ​
 onnx_path = "dirtyjudgment640320_pattern.onnx"
 torch.onnx.export(m, d, onnx_path)
 ​
 netron.start(onnx_path)

显示结果

两张图像中间的conv重复

 

pth转onnx

 def convert_model_to_ONNX(input_img_size, input_pth_model, output_ONNX):
     dummy_input = torch.randn(2, 3, input_img_size[1], input_img_size[0])
     model = SimpleNet()     #网络结构
     state_dict = torch.load(input_pth_model)
     new_state_dict = OrderedDict()
     for k, v in state_dict.items():
         name = k[7:]                    # remove `module.`
         new_state_dict[name] = v
     model.load_state_dict(new_state_dict)
     #model.load_state_dict(state_dict)
     input_names = ["input_image"]       #指定输入输出
     output_names = ["output_classification"]
     torch.onnx.export(model, dummy_input, output_ONNX, verbose=True, input_names=input_names,
                       output_names=output_names)

地址

GitHub - lutzroeder/netron: Visualizer for neural network, deep learning, and machine learning models