在做畢設的時候需要實現一個PyTorch原生代碼中沒有的并行算子，所以用到了這部分的知識，再不總結就要忘光了= =，本文內容主要是PyTorch的官方教程的各種傳送門，這些官方教程寫的都很好，以后就可以不用再浪費時間在百度上了。由于圖神經網絡計算框架PyG的代碼實現也是采用了擴展的方法，因此也可以當成下面總結PyG源碼文章的前導知識吧 。

第一種情況：使用PyThon擴展PyTorch

使用PyThon擴展PyTorch準確的來說是在PyTorch的Python前端實現自定義算子或者模型，不涉及底層C++的實現。這種擴展方式是所有擴展方式中最簡單的，也是官方首先推薦的，這是因為PyTorch在NVIDIA cuDNN，Intel MKL或NNPACK之類的庫的支持下已經對可能出現的CPU和GPU操作進行了高度優化，因此用Python擴展的代碼通常足夠快。

比如要擴展一個新的PyThon算子（torch.nn）只需要繼承torch.nn.Module并實現其forward方法即可。詳細的過程請參考官方教程傳送門：

Extending PyTorch?pytorch.org

第二種情況：使用pybind11構建共享庫形式的C++和CUDA擴展

但是如果我們想對代碼進行進一步優化，比如對自己的算子添加并行的CUDA實現或者連接個OpenCV的庫什么的，那么僅僅使用Python進行擴展就不能滿足需求；其次如果我們想序列化模型，在一個沒有Python環境的生產環境下部署，也需要我們使用C++重寫算法；最后考慮到考慮到多線程執行和性能原因，一般Python代碼也并不適合做部署。因此在對性能有要求或者需要序列化模型的場景下我們還是會用到C++擴展。

下面我先把官方教程傳送門放在這里：

CUSTOM C++ AND CUDA EXTENSIONS?pytorch.org

對于一種典型的擴展情況，比如我們要設計一個全新的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++ Frontend?pytorch.org

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

INSTALLING C++ DISTRIBUTIONS OF PYTORCH?pytorch.org

而Aten是ATen從根本上講是一個張量庫，在PyTorch中幾乎所有其他Python和C ++接口都在其上構建。它提供了一個核心Tensor類，在其上定義了數百種操作。這些操作大多數都具有CPU和GPU實現，Tensor該類將根據其類型向其動態調度。和Torch相比Aten更接近底層和核心邏輯。

Aten源代碼傳送門：

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

使用Aten聲明和操作張量的教程：

TENSOR BASICS?pytorch.org

由于Pyorch的C++后端文檔比較少，因此要多參考官方的例子，嘗試去模仿官方教程的代碼，同時可以通過Python前端的接口猜測后端接口的功能，如果沒有文檔了就讀一讀源碼，還是有不少注釋的，還能理解實現的邏輯。

第三種情況：為TORCHSCRIPT添加C++和CUDA擴展

首先簡單解釋一下TorchScript是什么，如果用官方的定義來說：“TorchScript是一種從PyTorch代碼創建可序列化和可優化模型的方法。任何TorchScript程序都可以從一個Python進程中保存并可以在一個沒有Python環境的進程中被加載??。”通俗來說TorchScript就是一個序列化模型（即Inference）的工具，它可以讓你的PyTorch代碼方便的在生產環境中部署，同時在將PyTorch代碼轉化TorchScript代碼時還會對你的模型進行一些性能上的優化。使用TorchScript完成模型的部署要比我們之前提到的使用C++重寫要簡單的多，因為是自動生成的。

TorchScript包含兩種序列化模型的方法：tracing和script，兩種方法各有其適用場景，由于和本文關系不大就不詳細展開了，具體的官方教程傳送門在此：

INTRODUCTION TO TORCHSCRIPT?pytorch.org

但是，TorchScript只能自動化的構造PyTorch的原生代碼，如果我們需要序列化自定義的C++擴展算子，則需要我們顯式的將這些自定義算子注冊到TorchScript中，所幸的是，這一過程其實非常簡單，整個過程和第二小節中使用pybind11構建共享庫的形式的C++和CUDA擴展十分類似。官方教程傳送門如下：

EXTENDING TORCHSCRIPT WITH CUSTOM C++ OPERATORS?pytorch.org

而對于自定義的C++類，如果要注冊到TorchScript要稍微復雜一些，官方教程傳送門如下：

EXTENDING TORCHSCRIPT WITH CUSTOM C++ CLASSES?pytorch.org

另外需要注意的是，如果想要編寫能夠被TorchScript編譯器理解的代碼，需要注意在C++自定義擴展算子參數中的數據類型，目前被TorchScript支持的參數數據類型有torch::Tensor，torch::Scalar(標量類型)，double，int64_t和std::vector，而像float，int，short這些是不能作為自定義擴展算子的參數數據類型的。

目前就先總結這么多吧，這點東西居然寫了一天，好累啊(* ￣︿￣)。

總結

以上是生活随笔為你收集整理的python interpreter 中没有torch_PyTorch扩展自定义PyThon/C++(CUDA)算子的若干方法总结的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。