关于Pytorch autograd机制反向传播计算梯度的几个tips

最近在用pytorch做性能分析的过程中，遇到了关于反向传播计算量的问题，趁此机会把pytorch的autograd机制看了一下，顺便来分享几个tips，帮助大家避坑。

本文的所有tips都是在阅读pytorch源码，包括torch/csrc/autograd目录下的autograd部分的C++源码总结出来的。

本文所有讨论的基础都建立在这个示例程序的调试之上，需要深入看一看Tensor类型关于梯度的几个属性：grad、grad_fn、next_functions

另外，请先阅读官方文档：The Fundamentals of Autograd — PyTorch Tutorials 2.2.1+cu121 documentation

示例程序，可以在pycharm调试一下这个程序，看一下loss变量的grad_fn、next_functions属性

import torch
BATCH_SIZE = 16
DIM_IN = 1000
HIDDEN_SIZE = 100
DIM_OUT = 10
​
class TinyModel(torch.nn.Module):
​
    def __init__(self):
        super(TinyModel, self).__init__()
​
        self.layer1 = torch.nn.Linear(DIM_IN, HIDDEN_SIZE)
        self.relu = torch.nn.ReLU()
        self.layer2 = torch.nn.Linear(HIDDEN_SIZE, DIM_OUT)
​
    def forward(self, x):
        x = self.layer1(x)
        x = self.relu(x)
        x = self.layer2(x)
        return x
​
some_input = torch.randn(BATCH_SIZE, DIM_IN, requires_grad=False)
ideal_output = torch.randn(BATCH_SIZE, DIM_OUT, requires_grad=False)
​
model = TinyModel()
​
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
​
prediction = model(some_input)
​
# loss_fn = torch.nn.CrossEntropyLoss()
​
# loss = loss_fn(prediction,ideal_output)
loss = (ideal_output - prediction).pow(2).sum()
print(loss)
​
​
loss.backward()
print(model.layer2.weight[0][0:10])
print(model.layer2.weight.grad[0][0:10])

进入调试，我们可以看到如下属性：

• 只有叶子节点的梯度被计算并保存（叶子节点=grad_fn=None）

• require_grad=True，表明参与计算图（track)

• AccumulateGrad grad_fn类型表示把所有batch_size计算出的梯度进行累加（不求平均，平均数因子早在之前的Divbackward已经算过了）：举例，广播的bias（原尺寸1 X 10，广播后16 X 10 ,batch_size=16）反向传播计算出的梯度尺寸也是16 X 10，bias为AccumulateGrad类型，则把梯度的第一个维度，即16进行累加，最终得到真正的梯度即 尺寸为1 X 10的梯度（与bias尺寸相同）

• Tbackward 即反转置，其实效果与普通的转置是一样的。一般用于计算权重W的梯度，因为pytorch autograd机制一般把权重W矩阵视为列主行矩阵进行计算（可以提高计算效率，参考源码torch/csrc/autograd/FunctionsManual.cpp中关于mm计算的几个函数），得到的结果与正确结果差一个转置，所以还需要进行一步Tbackward才能得到正确结果。

• Sumbackward 求和反向传播，就是把原梯度扩展到求和之前的尺寸，假如求和前的梯度尺寸 M X N ，求和后为标量，那么经过Sumbackward后，梯度尺寸还原为 M X N ，若在Sumbackward之前没有梯度被算出（即grad=None),那么会初始化元素梯度为全1矩阵

• Subbackward 减法反向传播，看情况，若被求梯度的元组在第二数位（即减号右边），则把梯度全部取负号即可。否则不变

• Powbackward 参考torch\csrc\autograd\FunctionsManual.cpp中的pow_backward函数

• 重要结论，pytorch在带有batch_size情况下的反向传播梯度计算中，计算量随batch_size倍增而倍增（基本上每一步反向计算量都要乘上batch_size因子），最后会统一把batch_size梯度累加起来，得到一个梯度。

• 在pytorch源码中我们一般找不到形如Subbackward函数的实现代码。这是因为pytorch采用一种torchgen的方式，通过重新组合一些原子操作（如add、mm)来生成复杂函数的C++代码。

• 详情请见pytorch/torchgen/api/autograd.py源码（自动生成反向传播算子）