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

编程算法

发布日期: 2021-03-24

作者: 韦阳

文章字数: 1.4k

阅读时长: 5 分

阅读次数:

在前面两篇教程中，我们详细讲解了如何编写cuda算子，并用PyTorch进行调用，并且详细讲述了三种编译cuda算子的方式，具体可以看前面两篇：

https://godweiyang.com/2021/03/18/torch-cpp-cuda/
https://godweiyang.com/2021/03/21/torch-cpp-cuda-2/

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

完整的代码还是放在了github仓库，欢迎大家star并fork：
https://github.com/godweiyang/torch-cuda-example

本文主要涉及到的是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类，我们可以用它来实现一下无法自动求导的操作，比如argmax这种不可导的函数。

我们需要实现两个函数，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
epoch 350: loss =    0.001
epoch 375: loss =    0.001
epoch 400: loss =    0.000
epoch 425: loss =    0.000
epoch 450: loss =    0.000
epoch 475: loss =    0.000