基于Netron库的PyTorch 2.0模型可视化

【图书推荐】《从零开始大模型开发与微调：基于PyTorch与ChatGLM》_《从零开始大模型开发与微调:基于pytorch与chatglm》-CSDN博客

前面章节带领读者完成了基于PyTorch 2.0的MNIST模型的设计，并基于此完成了MNIST手写体数字的识别。此时可能有读者对我们自己设计的模型结构感到好奇，如果能够可视化地显示模型结构就更好了。

读者可以自行在百度搜索Netron。Netron是一个深度学习模型可视化库，支持可视化地表示PyTorch 2.0的模型存档文件。因此，我们可以把3.1.2节中PyTorch的模型结构保存为文件，并通过Netron进行可视化展示。保存模型的代码如下：

class NeuralNetwork(torch.nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = torch.nn.Flatten()
        self.linear_relu_stack = torch.nn.Sequential(
            torch.nn.Linear(28*28,312),
            torch.nn.ReLU(),
            torch.nn.Linear(312, 256),
            torch.nn.ReLU(),
            torch.nn.Linear(256, 10)
        )
    def forward(self, input):
        x = self.flatten(input)
        logits = self.linear_relu_stack(x)

        return logits

#进行模型的保存
model = NeuralNetwork()
torch.save(model, './model.pth')		#将模型保存为pth文件

建议读者从GitHub上下载Netron，其主页提供了基于不同版本的安装方式，如图3-6所示。

读者可以依照操作系统的不同下载对应的文件，在这里安装的是基于Windows的.exe文件，安装后是一个图形界面，直接在界面上单击file操作符号打开我们刚才保存的.pth文件，显示结果如图3-7所示。

可以看到，此时我们定义的模型结构被可视化地展示出来了，每个模块的输入输出维度在图3-7上都展示出来了，单击深色部分可以看到每个模块更详细的说明，如图3-8所示。

本文节选自《从零开始大模型开发与微调：基于PyTorch与ChatGLM》，获得出版社和作者授权发布。

《从零开始大模型开发与微调：基于PyTorch与ChatGLM（人工智能技术丛书）》(王晓华)【摘要 书评 试读】- 京东图书 (jd.com)