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

原创

已于 2024-06-25 17:34:28 修改 · 763 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python

于 2024-06-25 08:12:07 首次发布

前面章节带领读者完成了基于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, 1