PyTorch中的C++扩展实现-创新互联

今天要聊聊用 PyTorch 进行 C++ 扩展。

在正式开始前，我们需要了解 PyTorch 如何自定义module。这其中，最常见的就是在 python 中继承torch.nn.Module，用 PyTorch 中已有的 operator 来组装成自己的模块。这种方式实现简单，但是，计算效率却未必最佳，另外，如果我们想实现的功能过于复杂，可能 PyTorch 中那些已有的函数也没法满足我们的要求。这时，用 C、C++、CUDA 来扩展 PyTorch 的模块就是最佳的选择了。

由于目前市面上大部分深度学习系统（TensorFlow、PyTorch 等）都是基于 C、C++ 构建的后端，因此这些系统基本都存在 C、C++ 的扩展接口。PyTorch 是基于 Torch 构建的，而 Torch 底层采用的是 C 语言，因此 PyTorch 天生就和 C 兼容，因此用 C 来扩展 PyTorch 并非难事。而随着 PyTorch2.0 的发布，官方已经开始考虑将 PyTorch 的底层代码用 caffe2 替换，因此他们也在逐步重构 ATen，后者是目前 PyTorch 使用的 C++ 扩展库。总的来说，C++ 是未来的趋势。至于 CUDA，这是几乎所有深度学习系统在构建之初就采用的工具，因此 CUDA 的扩展接口是标配。

本文用一个简单的例子，梳理一下进行 C++ 扩展的步骤，至于一些具体的实现，不做深入探讨。

PyTorch的C、C++、CUDA扩展

关于 PyTorch 的 C 扩展，可以参考官方教程或者这篇博文，其操作并不难，无非是借助原先 Torch 提供的<TH/TH.h>和等接口，再利用 PyTorch 中提供的torch.util.ffi模块进行扩展。需要注意的是，随着 PyTorch 版本升级，这种做法在新版本的 PyTorch 中可能会失效。

本文主要介绍 C++（未来可能加上 CUDA）的扩展方法。

C++扩展

首先，介绍一下基本流程。在 PyTorch 中扩展 C++/CUDA 主要分为几步：

1. 安装好 pybind11 模块（通过 pip 或者 conda 等安装），这个模块会负责 python 和 C++ 之间的绑定；
2. 用 C++ 写好自定义层的功能，包括前向传播forward和反向传播backward；
3. 写好 setup.py，并用 python 提供的setuptools来编译并加载 C++ 代码。
4. 编译安装，在 python 中调用 C++ 扩展接口。

接下来，我们就用一个简单的例子（z=2x+y）来演示这几个步骤。

第一步

安装 pybind11 比较简单，直接略过。我们先写好 C++ 相关的文件：

头文件 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);

另外有需要云服务器可以了解下创新互联scvps.cn，海内外云服务器15元起步，三天无理由+7*72小时售后在线，公司持有idc许可证，提供“云服务器、裸金属服务器、高防服务器、香港服务器、美国服务器、虚拟主机、免备案服务器”等云主机租用服务以及企业上云的综合解决方案，具有“安全稳定、简单易用、服务可用性高、性价比高”等特点与优势，专为企业上云打造定制，能够满足用户丰富、多元化的应用场景需求。