您的位置:首页 > 文旅 > 美景 > 基于Netron库的PyTorch 2.0模型可视化

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

2024/12/23 10:43:10 来源:https://blog.csdn.net/brucexia/article/details/139945267  浏览:    关键词:基于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)

版权声明:

本网仅为发布的内容提供存储空间,不对发表、转载的内容提供任何形式的保证。凡本网注明“来源:XXX网络”的作品,均转载自其它媒体,著作权归作者所有,商业转载请联系作者获得授权,非商业转载请注明出处。

我们尊重并感谢每一位作者,均已注明文章来源和作者。如因作品内容、版权或其它问题,请及时与我们联系,联系邮箱:809451989@qq.com,投稿邮箱:809451989@qq.com