在前面两篇教程中，我们详细讲解了如何编写cuda算子，并用PyTorch进行调用，并且详细讲述了三种编译cuda算子的方式，具体可以看前面两篇：
本文我们将讲解如何用自定义cuda算子搭建一个简单的神经网络，并实现反向传播，进行模型训练。
完整的代码还是放在了github仓库，欢迎大家star并fork：
https://github.com/godweiyang/torch-cuda-example
本文主要涉及到的是train.py这个代码，功能是搭建了一个PyTorch模型，并且调用了自定义的cuda算子，实现了自定义的反向传播函数，最终完成训练。
模型描述
之前我们实现了一个的tensor求和cuda算子，于是我们可以利用它来实现。
最终训练收敛后和都会趋近于0，模型没有输入，只有两个可训练的参数和。
搭建模型
首先我们还是像正常写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，也就是最终的损失函数对前向传播的返回值求导的结果。在我们这里的模型中，令
那么自定义cuda算子实现的就是这一步，而grad_output就是。我们自定义的cuda算子反向传播的导数就是和，然后根据链式求导法则就可以得到损失函数对每个参数的导数了。
反向传播返回值表示损失函数对前向传播每一个参数的梯度，所以个数必须等于前向传播除了ctx以外的其他参数个数，并且顺序也要一一对应。因为，所以返回值就是grad_output，grad_output和None，因为对常数不需要求导，所以直接返回空即可。
训练流程
最终训练流程和平常一样：
# 定义模型 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
小结
这三个教程暂时告一段落了，通过这些简单的例子，应该大致能学会如何自己写cuda算子，并且用PyTorch调用，完成模型训练了。
更复杂的模型其实基本的原理都是类似的，我不喜欢上来就讲解很复杂的大项目源码，我喜欢抽象出一个最简的example，这样更容易理解底层的原理，而不会被很多冗余的代码干扰。
