PyTorch扩展自定义PyThon/C++(CUDA)算子的若干方法总结

但是如果我们想对代码进行进一步优化，比如对自己的算子添加并行的CUDA实现或者连接个OpenCV的库什么的，那么仅仅使用Python进行扩展就不能满足需求；其次如果我们想序列化模型，在一个没有Python环境的生产环境下部署，也需要我们使用C++重写算法；最后考虑到考虑到多线程执行和性能原因，一般Python代码也并不适合做部署。因此在对性能有要求或者需要序列化模型的场景下我们还是会用到C++扩展。

下面我先把官方教程传送门放在这里：

CUSTOM C++ AND CUDA EXTENSIONSpytorch.org/tutorials/advanced/cpp_extension.html

对于一种典型的扩展情况，比如我们要设计一个全新的C++底层算子，其过程其实就三步：

第一步：使用C++编写算子的forward函数和backward函数

第二步：将该算子的forward函数和backward函数使用**pybind11**绑定到python上

第三步：使用setuptools/JIT/CMake编译打包C++工程为so文件

注意到在第一步中，我们不仅仅要实现forward函数也要实现backward函数，这是因为在C++端PyTorch目前不支持自动根据forward函数推导出backward函数，所以我们必须要对自己算子的反向传播过程完全清楚。一个需要注意的地方是，你可以选择直接在C++中继承torch::autograd类进行扩展；也可以像官方教程中那样在C++代码中实现forward和backward的核心过程，而在python端继承PyTorch的torch**.autograd.**Function类。

在C++端扩展forward函数和backward函数的需要注意以下规则：

（1）首先无论是forward函数还是backward函数都需要声明为静态函数。

（2）forward函数可以接受任意多的参数并且应该返回一个 variable list或者variable；forward函数需要将torch::autograd::AutogradContext 作为自己的第一个参数。Variables可以被使用ctx->save_for_backward保存，而其他数据类型可以使用ctx->saved_data以<std::string,at::IValue>pairs的形式保存在一个map中。

（3）backward函数第一个参数同样需要为torch::autograd::AutogradContext，其余的参数是一个variable_list，包含的变量数量与forward输出的变量数量相等。它应该返回和forward输入一样多的变量。保存在forward中的Variable变量可以通过ctx->get_saved_variables而其他的数据类型可以通过ctx->saved_data获取。

请注意，backward的输入参数是自动微分系统反传回来的参数梯度值，其需要和forward函数的返回值位置一一对应的；而backward的返回值是对各参数根据自动微分规则求导后的梯度值，其需要和forward函数的输入参数位置一一对应，对于不需要求导的参数也需要使用空Variable占位。

// PyG的C++扩展就选择的是直接继承PyTorch的C++端的torch::autograd类进行扩展
// 下面是PyG的一个ScatterSum算子的扩展示例
// 不用纠结这个算子的具体内容，对扩展的算子的结构有一个大致了解即可
class ScatterSum : public torch::autograd::Function<ScatterSum> {
public:// AutogradContext *ctx指针可以操作static variable_list forward(AutogradContext *ctx, Variable src,Variable index, int64_t dim,torch::optional<Variable> optional_out,torch::optional<int64_t> dim_size) {dim = dim < 0 ? src.dim() + dim : dim;ctx->saved_data["dim"] = dim;ctx->saved_data["src_shape"] = src.sizes();index = broadcast(index, src, dim);auto result = scatter_fw(src, index, dim, optional_out, dim_size, "sum");auto out = std::get<0>(result);ctx->save_for_backward({index});// 如果在扩展的C++代码中使用非Aten内建操作修改了tensor的值，需要对其进行脏标记if (optional_out.has_value())ctx->mark_dirty({optional_out.value()});  return {out};}// grad_outs是out参数反传回来的梯度值static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {auto grad_out = grad_outs[0];auto saved = ctx->get_saved_variables();auto index = saved[0];auto dim = ctx->saved_data["dim"].toInt();auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());auto grad_in = torch::gather(grad_out, dim, index, false);// 不需要求导的参数需要空Variable占位return {grad_in, Variable(), Variable(), Variable(), Variable()};}
};

由于涉及到在C++环境下操作张量和反向传播等操作，因此我们需要对PyTorch的C++后端的库有所了解，主要就是Torch和Aten这两个库，下面我简要介绍一下这两兄弟。

其中Torch是PyTorch的C++底层实现（PS：其实是先有的Torch后有的PyTorch，从名字也能看出来），FB在编码PyTorch的时候就有意将PyTorch的接口和Torch的接口设计的十分类似，因此如果你对PyTorch很熟悉的话那么你也会很快的对Torch上手。

Torch官方文档传送门：

The C++ Frontendpytorch.org/cppdocs/frontend.html

安装PyTorch的C++前端的官方教程：

INSTALLING C++ DISTRIBUTIONS OF PYTORCHpytorch.org/cppdocs/installing.html

而Aten是ATen从根本上讲是一个张量库，在PyTorch中几乎所有其他Python和C ++接口都在其上构建。它提供了一个核心Tensor类，在其上定义了数百种操作。这些操作大多数都具有CPU和GPU实现，Tensor该类将根据其类型向其动态调度。和Torch相比Aten更接近底层和核心逻辑。

Aten源代码传送门：

https://github.com/zdevito/ATen/tree/master/aten/srcgithub.com/zdevito/ATen/tree/master/aten/src

使用Aten声明和操作张量的教程：

TENSOR BASICSpytorch.org/cppdocs/notes/tensor_basics.html

由于Pyorch的C++后端文档比较少，因此要多参考官方的例子，尝试去模仿官方教程的代码，同时可以通过Python前端的接口猜测后端接口的功能，如果没有文档了就读一读源码，还是有不少注释的，还能理解实现的逻辑。