netron可视化模型，新手了解模型结构神器

需求描述：刚刚接触神经网络，因为对代码不敏感，可以选择可视化模型，把结构画出来，然后再对着代码看，方便理解。
方便你去了解模型的层之间的连接，输出和输入都是些什么。
方法：netron（可以有离线、在线、PC端软件），这里以 在线 和 离线 两种进行说明（在线：就是代码运行你就能看到模型结构，结束模型不能更新；离线：把模型结构保存下来，上传再查看）

第一步：安装netron

1. 在你的虚拟环境下安装netron

pip install netron

更多关于netron的细节可以通过netron项目的网页查看：
方式一： netron的GitHub项目网址： https://github.com/lutzroeder/netron
方式二： 访问不了GitHub看这个： https://gitcode.com/lutzroeder/netron/overview
其他系统安装办法：截图来于官网

2. 如果没有onnx，也需要安装

pip install onnx

注：其实这个我不确定，netron支持的格式这么多，不仅是onnx，所以用其他方法保存其他文件格式，应该也能识别，下面是netron支持的文件格式，来源官网：

第二步：修改实例化代码

class your_net(nn.Module):
	pass

if __name__ == "__main__":  
    model = your_net  
    input_size = torch.randn(1, 3, 256, 256)  
    out_size = model(input_size)  
    print(out_size.data.shape)  # 输出结果的size
    onnx_path = "./vis/vis_your_net.onnx"  # 模型onx格式的保存路径和命名
    torch.onnx.export(model, input_size, onnx_path)  
    netron.start(onnx_path)
    

以上可视化的一个例子，需要说明：

1. input_size这里是输出的一个模拟数组，模拟输入图像经过model时候的size等改变情况，根据你的模型修改input_size;
2. onnx_path这里后面的保存文件要加.onnx后缀名，不然后面看的模型会出问题（可以试一下不加，然后你对比一下区别，就明白了）
3. start我感觉就是传给可视化端口，默认的是8081

如果这个例子感觉不全，可以看到最后，有一个完整的代码，搭建了简单的神经网络，并进行实例化+可视化。

第三步：运行代码-进入浏览器可视化模型

我这里用的是UNet进行可视化，运行代码，可以出现转发端口，可以直接点击蓝色的那个网址就好，或者在浏览器输入： http://localhost:8081/

等待加载完毕，就可以看到模型了：

第四步：保存png文件 或 探索玩法（可选）

如图，我展示的是UNet，左边有一些选项，可以保存png格式，还有很多选项可以探索，浏览器端还允许查找的方式，找到你的结构（模型比较大很实用）

没有选择tensorboard主要是因为麻烦，版本兼容问题，其他的一些可视化方案也尝试过，不好搞，没有这个netron直观和方便。
我用过的可视化方法有：

• torchviz
• graphviz
• tensorboard
• netron

个人还是觉得netron比较好用

离线其实就是把保存的.onnx模型数据上传到网页版的netron（网址： https://netron.app/）

第五步：你可以使用的例子：

import torch  
import torch.nn as nn  
import netron  
from torchsummary import summary  
  
class SimpleCNN(nn.Module):  
    def __init__(self, in_channels, num_classes):  
        super(SimpleCNN, self).__init__()  
        self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1)  
        self.relu1 = nn.ReLU()  
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  
  
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)  
        self.relu2 = nn.ReLU()  
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)  
  
        self.fc1 = nn.Linear(128 * 64 * 64, 512)  
        self.relu3 = nn.ReLU()  
        self.fc2 = nn.Linear(512, num_classes)  
  
    def forward(self, x):  
        x = self.conv1(x)  
        x = self.relu1(x)  
        x = self.pool1(x)  
  
        x = self.conv2(x)  
        x = self.relu2(x)  
        x = self.pool2(x)  
  
        x = x.view(x.size(0), -1)  # Flatten the output for fully connected layers  
        x = self.fc1(x)  
        x = self.relu3(x)  
        x = self.fc2(x)  
  
        return x  
  
  
if __name__ == "__main__":  
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
    # 定义模型  
    model = SimpleCNN(in_channels=3, num_classes=2).to(device)  
    summary(model, input_size=(3, 256, 256))  
  
    # 随机生成输入数据  
    input_size = torch.randn(1, 3, 256, 256).to(device)  
    # 前向传播  
    output = model(input_size)  
    # 打印输出形状  
    print(output.shape)  
  
    onnx_path = "./vis/SimpleCNN.onnx"  # 模型onx格式的保存路径和命名  
    torch.onnx.export(model, input_size, onnx_path)  
    netron.start(onnx_path)

注意：我把数据和模型都移到了GPU

然后运行就可以看到这里出现这样的结果：

相应的8081端口输出，浏览器可以查看模型如下：

恭喜你！你又学会一招！
感谢开发netron的大佬！