【Pytorch学习笔记】2.动手生成计算图——将Tensor间的计算流程和梯度传递可视化，使用torchviz生成计算图

在学习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/