三分钟教你如何PyTorch自定义反向传播-白红宇

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

阅读量：728 次

发布时间：2019-03-21

本文共 3434 字，大约阅读时间需要 11 分钟。

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

本文我们将讲解如何用自定义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.658epoch  25: loss =  727.122epoch  50: loss =  264.796epoch  75: loss =   96.431epoch 100: loss =   35.117epoch 125: loss =   12.789epoch 150: loss =    4.657epoch 175: loss =    1.696epoch 200: loss =    0.618epoch 225: loss =    0.225epoch 250: loss =    0.082epoch 275: loss =    0.030epoch 300: loss =    0.011epoch 325: loss =    0.004epoch 350: loss =    0.001epoch 375: loss =    0.001epoch 400: loss =    0.000epoch 425: loss =    0.000epoch 450: loss =    0.000epoch 475: loss =    0.000

小结

这三个教程暂时告一段落了，通过这些简单的例子，应该大致能学会如何自己写cuda算子，并且用PyTorch调用，完成模型训练了。

更复杂的模型其实基本的原理都是类似的，我不喜欢上来就讲解很复杂的大项目源码，我喜欢抽象出一个最简的example，这样更容易理解底层的原理，而不会被很多冗余的代码干扰。

转载地址：http://qjaez.baihongyu.com/

你可能感兴趣的文章