三分钟教你如何PyTorch自定义反向传播

本文我们将讲解如何用自定义cuda算子搭建一个简单的神经网络，并实现反向传播，进行模型训练。

完整的代码还是放在了github仓库，欢迎大家star并fork：

本文主要涉及到的是 train.py 这个代码，功能是搭建了一个PyTorch模型，并且调用了自定义的cuda算子，实现了自定义的反向传播函数，最终完成训练。

模型描述

之前我们实现了一个 a + b 的tensor求和cuda算子，于是我们可以利用它来实现 \mathcal{L} = a^2 + b^2 。

最终训练收敛后 a 和 b 都会趋近于0，模型没有输入，只有两个可训练的参数 a 和 b 。

搭建模型

首先我们还是像正常写PyTorch模型那样搭建一个模型，代码如下：

class AddModel(nn.Module):
    def __init__(self, n):
        super(AddModel, self).__init__()
        # tensor长度
        self.n = n
        # 定义可训练参数a和b
        self.a = nn.Parameter(torch.Tensor(self.n))
        self.b = nn.Parameter(torch.Tensor(self.n))
        # 正态分布初始化参数a和b
        self.a.data.normal_(mean=0.0, std=1.0)
        self.b.data.normal_(mean=0.0, std=1.0)
    def forward(self):
        # 求a^2与b^2
        a2 = torch.square(self.a)
        b2 = torch.square(self.b)
        # 调用自定义cuda算子对两个平方数求和
        c = AddModelFunction.apply(a2, b2, self.n)
        return c

重点就在调用自定义cuda算子那一行 AddModelFunction.apply() ，你也可以写成 c = a2 + b2 。不过这里我们为了演示如何使用自定义cuda算子，所以不这么干了。

实现自定义cuda算子前向和反向传播

下面就是如何实现 AddModelFunction.apply() 函数了，我们先来看一下具体代码：

class AddModelFunction(Function):
    @staticmethod
    def forward(ctx, a, b, n):
        c = torch.empty(n).to(device="cuda:0")
        if args.compiler == 'jit':
            cuda_module.torch_launch_add2(c, a, b, n)
        elif args.compiler == 'setup':
            add2.torch_launch_add2(c, a, b, n)
        elif args.compiler == 'cmake':
            torch.ops.add2.torch_launch_add2(c, a, b, n)
        else:
            raise Exception("Type of cuda compiler must be one of jit/setup/cmake.")
        return c
    @staticmethod
    def backward(ctx, grad_output):
        return (grad_output, grad_output, None)

这个类继承的是 torch.autograd.Function 类，我们可以用它来实现一下无法自动求导的操作，比如 arxmax 这种不可导的函数。

我们需要实现两个函数， forward 和 backward ，分别用来前向和反向传播，注意都得声明成静态函数。

前向传播接收多个参数，第一个固定为 ctx ，用来存储反向传播中可能会用到的一些上下文，比如 input 和一些前向过程中的中间变量等等，其他参数随你定。然后我们根据上一教程中调用cuda算子的方法计算得到求和结果，进行返回。

反向传播接收两个参数，第一个同样是 ctx ，里面存着前向过程中保存的一些上下文变量信息。第二个是 grad_output ，也就是最终的损失函数对前向传播的返回值求导的结果。在我们这里的模型中，令

a2 = a^2, b2 = b^2, s = a2 + b2, \mathcal{L} = s \\

那么自定义cuda算子实现的就是 s = a2 + b2 这一步，而 grad_output 就是 \frac{\partial \mathcal{L}}{\partial s} 。我们自定义的cuda算子反向传播的导数就是 \frac{\partial s}{\partial a2} 和 \frac{\partial s}{\partial b2} ，然后根据链式求导法则就可以得到损失函数对每个参数的导数了。

反向传播返回值表示损失函数对前向传播每一个参数的梯度，所以个数必须等于前向传播除了 ctx 以外的其他参数个数，并且顺序也要一一对应。因为 \frac{\partial s}{\partial a2} = \frac{\partial s}{\partial b2} = 1 ，所以返回值就是 grad_output ， grad_output 和 None ，因为对常数 n 不需要求导，所以直接返回空即可。

训练流程

最终训练流程和平常一样：

# 定义模型
model = AddModel(n)
# 将模型中所有参数拷贝到GPU端
model.to(device="cuda:0")
# 定义优化器
opt = torch.optim.SGD(model.parameters(), lr=0.01)
for epoch in range(500):
    # 清空优化器缓存
    opt.zero_grad()
    # 前向传播
    output = model()
    # 求loss
    loss = output.sum()
    # 反向传播
    loss.backward()
    # 更新参数
    opt.step()
    if epoch % 25 == 0:
        print("epoch {:>3d}: loss = {:>8.3f}".format(epoch, loss))

最终损失函数降到了0，log信息如下：

Loading extension module add2...
Initializing model...
Initializing optimizer...
Begin training...
epoch   0: loss = 1996.658
epoch  25: loss =  727.122
epoch  50: loss =  264.796
epoch  75: loss =   96.431
epoch 100: loss =   35.117
epoch 125: loss =   12.789
epoch 150: loss =    4.657
epoch 175: loss =    1.696
epoch 200: loss =    0.618
epoch 225: loss =    0.225
epoch 250: loss =    0.082
epoch 275: loss =    0.030
epoch 300: loss =    0.011
epoch 325: loss =    0.004