今天要聊聊用 PyTorch 進行 C++ 擴展。

十年建站經(jīng)驗， 做網(wǎng)站、網(wǎng)站制作客戶的見證與正確選擇。創(chuàng)新互聯(lián)提供完善的營銷型網(wǎng)頁建站明細報價表。后期開發(fā)更加便捷高效，我們致力于追求更美、更快、更規(guī)范。

在正式開始前，我們需要了解 PyTorch 如何自定義module。這其中，最常見的就是在 python 中繼承torch.nn.Module，用 PyTorch 中已有的 operator 來組裝成自己的模塊。這種方式實現(xiàn)簡單，但是，計算效率卻未必最佳，另外，如果我們想實現(xiàn)的功能過于復(fù)雜，可能 PyTorch 中那些已有的函數(shù)也沒法滿足我們的要求。這時，用 C、C++、CUDA 來擴展 PyTorch 的模塊就是最佳的選擇了。

由于目前市面上大部分深度學(xué)習(xí)系統(tǒng)（TensorFlow、PyTorch 等）都是基于 C、C++ 構(gòu)建的后端，因此這些系統(tǒng)基本都存在 C、C++ 的擴展接口。PyTorch 是基于 Torch 構(gòu)建的，而 Torch 底層采用的是 C 語言，因此 PyTorch 天生就和 C 兼容，因此用 C 來擴展 PyTorch 并非難事。而隨著 PyTorch2.0 的發(fā)布，官方已經(jīng)開始考慮將 PyTorch 的底層代碼用 caffe2 替換，因此他們也在逐步重構(gòu) ATen，后者是目前 PyTorch 使用的 C++ 擴展庫?？偟膩碚f，C++ 是未來的趨勢。至于 CUDA，這是幾乎所有深度學(xué)習(xí)系統(tǒng)在構(gòu)建之初就采用的工具，因此 CUDA 的擴展接口是標配。

本文用一個簡單的例子，梳理一下進行 C++ 擴展的步驟，至于一些具體的實現(xiàn)，不做深入探討。

PyTorch的C、C++、CUDA擴展

關(guān)于 PyTorch 的 C 擴展，可以參考官方教程或者這篇博文，其操作并不難，無非是借助原先 Torch 提供的<TH/TH.h>和等接口，再利用 PyTorch 中提供的torch.util.ffi模塊進行擴展。需要注意的是，隨著 PyTorch 版本升級，這種做法在新版本的 PyTorch 中可能會失效。

本文主要介紹 C++（未來可能加上 CUDA）的擴展方法。

C++擴展

首先，介紹一下基本流程。在 PyTorch 中擴展 C++/CUDA 主要分為幾步：

1. 安裝好 pybind11 模塊（通過 pip 或者 conda 等安裝），這個模塊會負責(zé) python 和 C++ 之間的綁定；
2. 用 C++ 寫好自定義層的功能，包括前向傳播forward和反向傳播backward；
3. 寫好 setup.py，并用 python 提供的setuptools來編譯并加載 C++ 代碼。
4. 編譯安裝，在 python 中調(diào)用 C++ 擴展接口。

接下來，我們就用一個簡單的例子（z=2x+y）來演示這幾個步驟。

第一步

安裝 pybind11 比較簡單，直接略過。我們先寫好 C++ 相關(guān)的文件：

頭文件 test.h

#include 
#include 

// 前向傳播
torch::Tensor Test_forward_cpu(const torch::Tensor& inputA,
              const torch::Tensor& inputB);
// 反向傳播
std::vector Test_backward_cpu(const torch::Tensor& gradOutput);

另外有需要云服務(wù)器可以了解下創(chuàng)新互聯(lián)scvps.cn，海內(nèi)外云服務(wù)器15元起步，三天無理由+7*72小時售后在線，公司持有idc許可證，提供“云服務(wù)器、裸金屬服務(wù)器、高防服務(wù)器、香港服務(wù)器、美國服務(wù)器、虛擬主機、免備案服務(wù)器”等云主機租用服務(wù)以及企業(yè)上云的綜合解決方案，具有“安全穩(wěn)定、簡單易用、服務(wù)可用性高、性價比高”等特點與優(yōu)勢，專為企業(yè)上云打造定制，能夠滿足用戶豐富、多元化的應(yīng)用場景需求。

本文名稱：PyTorch中的C++擴展實現(xiàn)-創(chuàng)新互聯(lián)
文章轉(zhuǎn)載：http://weahome.cn/article/djcpcc.html