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前端的接口猜測后端接口的功能，如果沒有文檔了就讀一讀源碼，還是有不少注釋的，還能理解實現的邏輯。