学术, 机器学习

使用torchviz实现PyTorch计算图的可视化

本篇使用 torchviz 软件包实现 PyTorch 计算图的可视化。

软件包安装(GitHub 网址为:https://github.com/szagoruyko/pytorchviz):

pip install torchviz

另外,可能还需要下载安装:https://graphviz.org/download/,并添加到系统的环境变量中,不然运行时可能会报错。

代码例子:

"""
This code is supported by the website: https://www.guanjihuan.com
The newest version of this code is on the web page: https://www.guanjihuan.com/archives/40353
"""

import torch
import torchviz


# 简单网络的例子
class SimpleNet(torch.nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = torch.nn.Linear(10, 6)
        self.fc2 = torch.nn.Linear(6, 2)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x
model = SimpleNet()
input = torch.randn(1, 10)
output = model(input)
graph = torchviz.make_dot(output, params=dict(model.named_parameters()))
graph.render("Simple_net_graph") # 保存计算图为 PDF 文件


# 卷积网络的例子
class ConvolutionalNeuralNetwork(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.convolutional_layer_1 = torch.nn.Conv2d(in_channels=1, out_channels=10, kernel_size=3, stride=1, padding=1)
        self.convolutional_layer_2 = torch.nn.Conv2d(in_channels=10, out_channels=10, kernel_size=3, stride=1, padding=1)
        self.pooling_layer = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        self.hidden_layer_1 = torch.nn.Linear(in_features=10*7*7, out_features=10)
        self.hidden_layer_2 = torch.nn.Linear(in_features=10, out_features=10)
        self.output_layer = torch.nn.Linear(in_features=10, out_features=1)
    def forward(self, x):
        channel_output_1 = torch.nn.functional.relu(self.pooling_layer(self.convolutional_layer_1(x)))
        channel_output_2 = torch.nn.functional.relu(self.pooling_layer(self.convolutional_layer_2(channel_output_1)))
        channel_output_2 = torch.flatten(channel_output_2, 1)
        hidden_output_1 = torch.nn.functional.relu(self.hidden_layer_1(channel_output_2))
        hidden_output_2 = torch.nn.functional.relu(self.hidden_layer_2(hidden_output_1))
        output = self.output_layer(hidden_output_2)
        return output
model = ConvolutionalNeuralNetwork()
input = torch.randn(15, 1, 28, 28)
output = model(input)
graph = torchviz.make_dot(output, params=dict(model.named_parameters()))
graph.render("CNN_graph") # 保存计算图为 PDF 文件

运行结果:

总结:

  • 对于线性层,节点是AddmmBackward0。其中,bias之后是AccumulateGrad,weight之后是AccumulateGrad和TBackward0。
  • 对于卷积层,节点是ConvolutionBackward0。其中,weight和bias之后都是AccumulateGrad。
  • 一般来说,看代码就可以直接想象出计算图的大概样子。这里生成的图只是用于辅助,方便直观地看出神经网络的复杂程度。

图中节点的说明(以下由 AI 生成,使用的模型为 ChatGPT ):

  • AccumulateGrad: 这个节点表示梯度的累积。在反向传播过程中,梯度是被累积的,因为在每一次反向传播调用中,梯度都会被计算并加到之前的梯度上,以便进行梯度下降更新。
  • TBackward0: 这个节点通常代表Tensor的反向传播操作。当在Tensor上进行某些操作并计算梯度时,这个节点将会出现,表示该Tensor的反向传播。
  • AddmmBackward0: 这是矩阵乘法操作的反向传播。在神经网络中,矩阵乘法常用于线性层(fully connected layer)的计算,而这个节点则表示反向传播的计算。
  • ConvolutionBackward0: 这是卷积操作的反向传播。在卷积神经网络中,卷积操作是一种常见的操作,这个节点表示卷积层的反向传播计算。
  • MaxPool2DWithIndicesBackward0: 这个节点是最大池化操作的反向传播。在最大池化层中,通过取池化窗口中的最大值来减小输入的空间大小,这个节点表示最大池化操作的反向传播。
  • ReluBackward0: 这是ReLU激活函数的反向传播。ReLU函数是一种常用的非线性激活函数,在深度学习中广泛使用。这个节点表示ReLU函数的反向传播计算。
  • ReshapeAliasBackward0: 这个节点是reshape操作的反向传播。当对张量进行reshape操作时,这个节点表示reshape操作的反向传播计算。
108 次浏览

【说明:本站主要是个人的一些笔记和代码分享,内容可能会不定期修改。为了使全网显示的始终是最新版本,这里的文章未经同意请勿转载。引用请注明出处:https://www.guanjihuan.com

发表评论

您的电子邮箱地址不会被公开。 必填项已用 * 标注

Captcha Code