PyTorch 2.0深度学习从零开始学最新章节_王晓华著

3.2　PyTorch 2.0模型结构输出与可视化

上一节中完成了基于PyTorch 2.0的MNIST模型的设计，并完成了MNIST手写体的识别。此时，可能会有读者对我们自己设计的模型结构感到好奇，如果有一种能够可视化模型结构的现成方法，用起来就非常方便了。

3.2.1　查看模型结构和参数信息

为了解决模型结构的展示问题，PyTorch官方提供了对应的建议模型打印工具，即直接调用print函数来完成，例如对上一节中我们实现的MNIST模型：

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

读者可以直接使用如下函数完成模型参数的打印：

     if __name__ == '__main__':
         model = NeuralNetwork()
         print(model)

打印结果如图3-5所示。

图3-5　对模型具体使用的函数及其对应的参数进行打印

可以看到此结果是对模型具体使用的函数及其对应的参数进行打印。

为了更进一步简化对模型参数的打印，读者可以使用作者提供的对参数和结构进行打印的函数，代码如下所示：

     params = list(model.parameters())
     k = 0
     for i in params:
             l = 1
             print("该层的结构：" + str(list(i.size())))
             for j in i.size():
                 l *= j
             print("该层参数和：" + str(l))
             k = k + l
     print("总参数数量和：" + str(k))

运行完此段代码后，可以对每层的节点输出和参数总量进行打印，结果如图3-6所示。

图3-6　对每层的节点输出和参数总量进行打印

3.2.2　基于netron库的PyTorch 2.0模型可视化

上一小节讲解了模型结构的输出，但是相对于简单的文本化输出方式，PyTorch第三方提供了许多具有直观表示的模型可视化展示方式。netron就是一个深度学习模型可视化库，支持可视化表示PyTorch 2.0的模型存档文件。我们可以将上一小节中PyTorch的模型结构进行保存，并通过netron进行可视化展示。模型保存代码如下所示：

     import torch
     device = "cuda"   #在这里默认使用GPU，如果出现运行问题可以将其改成cpu模式
　　　　　
     #设定的多层感知机网络模型
     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文件

读者可以自行百度netron的下载地址，建议从GitHub上下载netron，其主页上也提供了不同版本的安装方式，如图3-7所示。