在前面两篇教程中，我们详细讲解了如何编写cuda算子，并用PyTorch进行调用，并且详细讲述了三种编译cuda算子的方式，具体可以看前面两篇：
PyTorch自定义CUDA算子教程与运行时间分析
详解PyTorch编译并调用自定义CUDA算子的三种方式
本文我们将讲解如何用自定义cuda算子搭建一个简单的神经网络，并实现反向传播，进行模型训练。
完整的代码还是放在了github仓库，欢迎大家star并fork：
https://github.com/godweiyang/torch-cuda-example
本文主要涉及到的是train.py这个代码，功能是搭建了一个PyTorch模型，并且调用了自定义的cuda算子，实现了自定义的反向传播函数，最终完成训练。
模型描述
之前我们实现了一个
的tensor求和cuda算子，于是我们可以利用它来实现
。
最终训练收敛后
和
都会趋近于0，模型没有输入，只有两个可训练的参数
和
。
搭建模型
首先我们还是像正常写PyTorch模型那样搭建一个模型，代码如下：
代码语言：javascript复制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，也就是最终的损失函数对前向传播的返回值求导的结果。在我们这里的模型中，令
那么自定义cuda算子实现的就是
这一步，而grad_output就是
。我们自定义的cuda算子反向传播的导数就是
和
，然后根据链式求导法则就可以得到损失函数对每个参数的导数了。
反向传播返回值表示损失函数对前向传播每一个参数的梯度，所以个数必须等于前向传播除了ctx以外的其他参数个数，并且顺序也要一一对应。因为
，所以返回值就是grad_output，grad_output和None，因为对常数
不需要求导，所以直接返回空即可。
训练流程
最终训练流程和平常一样：
代码语言：javascript复制# 定义模型
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信息如下：
代码语言：javascript复制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
小结
这三个教程暂时告一段落了，通过这些简单的例子，应该大致能学会如何自己写cuda算子，并且用PyTorch调用，完成模型训练了。
更复杂的模型其实基本的原理都是类似的，我不喜欢上来就讲解很复杂的大项目源码，我喜欢抽象出一个最简的example，这样更容易理解底层的原理，而不会被很多冗余的代码干扰。
- END -
我是godweiyang，华东师范大学计算机系本硕专业第一，字节跳动AI Lab NLP算法工程师，秋招斩获上海三家互联网大厂ssp offer，主要研究方向为机器翻译、句法分析、模型压缩与加速。最大特点就是脾气好、有耐心，有任何问题都可以随时咨询我，不管是技术上的还是生活上的。
