PyTorch实战：借助torchviz可视化计算图与梯度传递

文章目录

• • Tensor计算的可视化（线性回归为例）
    • 如何使用可视化库torchviz
    • • 安装graphviz软件
        • 安装torchviz库
        • 使用 torchviz.make_dot()

在学习Tensor时，将张量y用张量x表示，它们背后会有一个函数表达关系，y的
grad_fn
会被赋予一个对应的函数。先定义的x是一个叶子节点，将所有Tensor节点的计算连接起来就可以用一个
有向无环图
(DAG)来表示，称为
计算图
（Computational graphs）。

Computational graphs 例图：

有了图之后就可以清晰直观地理解这个模型的计算过程(forward)和梯度传递(backward)。

在初学线性回归模型的时候，同样可以把手写的线性回归模型以计算图的形式可视化表示出来，可便于深入理解代码背后的计算过程。

Tensor计算的可视化（线性回归为例）

使用的同样是“动手学深度学习的
线性回归从零实现
的例子”。

该模型的定义函数和损失函数为：

def linreg(X, w, b):
    return torch.mm(X, w) + b
    
def squared_loss(y_hat, y): 
    return (y_hat - y.view(y_hat.size())) ** 2 / 2

在训练模型时定义的损失函数为：

net = linreg
loss = squared_loss
# 因为loss完是一个10*1的张量，所以需要sum一下转换成标量以便标量求导
l = loss(net(X, w, b), y).sum()

我们可以将损失函数 l 的计算图表示出来，以
理解梯度是怎么传播
的。

这里使用torchviz的make_dot()函数将这些计算节点表示出来，在
例子
末尾添加代码，执行：

data_iter1 = data_iter(batch_size, features, labels)
X, y = next(data_iter1)  # 另取一个batch数据

l = loss(net(X, w, b), y).sum()

# 判断各节点是否是叶子节点
print(X.is_leaf) 
print(y.is_leaf)
print(w.is_leaf)
print(b.is_leaf)
print(l.is_leaf)
print(X.grad_fn)
print(l.grad_fn)

from torchviz import make_dot
make_dot(l.mean())
# 这里添加mean()是对之前的取sum()求平均，具体可以参考 自动求导的实现 相关知识。

结果显示：

True

True

True

True

False

None

<SumBackward0 object at 0x000001DC230C7DA0>

下面可以对该线性回归模型的损失函数进行分析
：

可以看到X,y,w,b都是计算图中的叶子节点，l表示计算流终点，它不是叶子节点。

X是叶子节点，没有
grad_fn
，即没有创建该Tensor的Function。

l代表父节点，它有创建该Tensor的Function，可以看到它的
grad_fn
是
SumBackward0
，也就是图中倒数第三个节点。

图中可以清晰地看到计算过程和梯度传递的可视化，包含了计算图的信息。左上蓝色框是shape为(2,1)的权重w，右上是b。因为X和y没有被赋予梯度，所以不出现在图中。

箭头的方向是计算执行的朝向，每次计算朝向下一个带
grad_fn
的节点，梯度则根据这个流向进行反传。

以上是最简单的线性模型，在之后的神经网络模型中，同样可以使用类似的方法对
模型
或者
损失函数
等计算过程进行可视化分析。可视化工具torchviz的安装使用见下。

如何使用可视化库torchviz

可视化需要安装torchviz库和graphviz软件。

安装graphviz软件

1. 下载：

https://graphviz.org/download/

选择对应的平台的安装方式，Windows是下载安装包。

下载好以后根据向导安装，注意根据提示把graphviz添加到环境变量中去。
2. 安装成功后再cmd界面输入
dot -version
可显示版本信息

3. 安装python库：

pip install graphviz

安装torchviz库

pip install torchviz

使用 torchviz.make_dot()

使用 torchviz.make_dot() 函数就可以将Tensor计算和梯度传播过程可视化了。推荐在jupyter环境下使用。

使用方法：直接在make_dot()中传入待分析的Tensor变量即可，比如上面的例子。

如果想自行控制计算图的展示，在jupyter notebook中使用display即可：

dot = torchviz.make_dot(l)  # make_dot返回一个dot（一个Diagraph对象）
display(dot)

其它进阶使用方法可以参考官方的文档：
地址
。

参考文献：

https://pytorch.org/blog/computational-graphs-constructed-in-pytorch/