pytorch的C++ extension写法

2020 年 4 月 17 日 极市平台

加入极市专业CV交流群，与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流！

同时提供每月大咖直播分享、真实项目需求对接、干货资讯汇总，行业技术交流。关注 极市平台 公众号，回复 加群，立刻申请入群~

本文授权转自知乎作者Monstarrrr，https://zhuanlan.zhihu.com/p/100459760。未经作者许可，不可二次转载。

2019年的最后一天，终于填了一个早就想了解的坑。就是关于pytorch如何自定义一个扩展，这里主要是说C++扩展。

首先为什么需要扩展？python调用C++的库也是可行的啊。刚开始我也在思考这个问题，觉得没有必要。但是后来深入了解了以后发现还是有必要的。举个栗子，调用始终是使用的是别人的东西，但是扩展则是通过他人的帮助来完成一个属于自己的东西。

pytorch的C++ extension和python的c/c++ extension其实原理差不多，本质上都是为了扩展各自的功能，当然也为了使程序运行更加有效率，差别在于pytorch的C++ extension实施步骤较python的c/c++ extension的要简化一些。

这里以实现神经网络自定义的layer为例：

先说一下基本的流程：

利用C++写好自定义层发功能，主要包括前向传播和方向传播，以及pybind11的内容。
写好setup.py脚本，并利用python提供的setuptools来编译并加载C++代码。
编译安装，在python中调用C++扩展接口

pybind11是python的一个库，主要负责python与C++11之间的通信

下面就以一个最简单的z=2x+y来看看如何一步步完成这样一个简单运算的layer。

第一步：编写头文件，这里就叫做test.h

/*test.h*/
#include <torch/extension.h>#include <vector>
// forward propagation
torch::Tensor Test_forward_cpu(const torch::Tensor& inputA, const torch::Tensor& inputB);
// backward propagation
std::vector<torch::Tensor> Test_backward_cpu(const torch::Tensor& gradOutput);

这里包含一个重要的头文件<torch/extension.h>

这个头文件里面包含很多重要的模块。如用于python和C++11交互的pybind11，以及包含Tensor的一系列定义操作，因为pytorch的基本数据单元是Tensor。

头文件写完以后就要开始写源文件了test.cpp

/*test.cpp*/
#include "test.h"
// part1:forward propagation
torch::Tensor Test_forward_cpu(const torch::Tensor& x, const torch::Tensor& y)
{
    AT_ASSERTM(x.sizes() == y.sizes());
    torch::Tensor z = torch::zeros(x.sizes());
    z = 2 * x + y;
    return z;
}

//part2:backward propagation
std::vector<torch::Tensor> Test_backward_cpu(const torch::Tensor& gradOutput)
{
    torch::Tensor gradOutputX = 2 * gradOutput * torch::ones(gradOutput.sizes());
    torch::Tensor gradOutputY = gradOutput * torch::ones(gradOutput.sizes());
    return {gradOutputX, gradOutputY};
}

// part3:pybind11 （将python与C++11进行绑定， 注意这里的forward，backward名称就是后来在python中可以引用的方法名）
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    m.def("forward", &Test_forward_cpu, "Test forward");
    m.def("backward", &Test_backward_cpu, "Test backward");
}

源文件cpp里面包含了三个部分，第一个部分是forward函数，第二个部分是backward函数，第三个部分是pytorch和C++交互的部分。

至此C++部分的工作就完成了。也就是我们的步骤一：利用C++写好自定义层发功能，主要包括前向传播和方向传播，以及pybind11的内容。

下面的工作就是pytorch如何识别和使用这个扩展程序了。

第二步：编写setup.py，这个文件的主要作用是用来编译C++文件以及建立链接关系。

现在的文件目录排布为：

setup.py中的内容为：

from setuptools import setup
import os
import glob
from torch.utils.cpp_extension import BuildExtension, CppExtension

# 头文件目录
include_dirs = os.path.dirname(os.path.abspath(__file__))
#源代码目录
source_file = glob.glob(os.path.join(working_dirs, 'src', '*.cpp'))

setup(
    name='test_cpp',  # 模块名称
    ext_modules=[CppExtension('test_cpp', sources=source_file, include_dirs=[include_dirs])],
    cmdclass={
        'build_ext': BuildExtension
    }
)

这一部分基本上算是一个固定的格式针对不同的问题需要修改的地方就是ext_modules参数，这里面根据实际的需要列表中可以存在多个CppExtension模块，也就是说可以同时编译多个C++文件。

例如像这样：

完成setup.py以后，需要在终端执行python setup.py install

NOTE:建议将扩展安装在个人虚拟环境中

这一步其实是包含了build+install执行的是先编译链接动态链接库，然后将构建好的文件以package的形式安装存放再当前开发环境的package的集中存放处，这样就相当于生成了一个完整的package了。和其他的如numpy，torch这些package没什么两样。

执行完这一步后就生成了这一堆东西：

这样，我们的第二步“写好setup.py脚本，并利用python提供的setuptools来编译并加载C++代码。”也完成了。

NOTE：此时如果在python的控制台中输入import test_cpp会得到这样的错误：

undefined symbol: _ZTIN3c1021AutogradMetaInterfaceE

原因是因为它还没有封装起来，暂时还见不得人~。

下面是最后一步：封装调用这个扩展（extension），先在与setup.py相同的目录下新建一个test.py

内容为：

from torch.autograd import Function
import torch
import test_cpp


class _TestFunction(Function):
    @staticmethod
    def forward(ctx, x, y):
        """
        It must accept a context ctx as the first argument, followed by any
        number of arguments (tensors or other types).
        The context can be used to store tensors that can be then retrieved
        during the backward pass."""
        return test_cpp.forward(x, y)

    @staticmethod
    def backward(ctx, gradOutput):
        gradX, gradY = test_cpp.backward(gradOutput)
        return gradX, gradY

# 封装成一个模块（Module）
class Test(torch.nn.Module):
    def __init__(self):
        super(Test, self).__init__()

    def forward(self, inputA, inputB):
        return _TestFunction.apply(inputA, inputB)

这是pytorch的autograd中的一个扩展函数的接口模板。基本pytorch中所有的层的前向传播和反向传播都是这样写的。

关于pytorch的方向传播的细节，有两个需要注意的点。其一，forward函数中有一个ctx变量，这是一定需要的。因为这里面会存一些对方向传播有用的变量（因为有些函数求导是需要用到前向计算过程中的一些计算结果）。backward中也有ctx参数，可以获取从forward函数中所保存的变量。第二个需要注意的点就是backward输出的都是关于变量的梯度，其数目要和forwad中输入的一致，这是一种强制性的要求，如果有些变量不需要求导，就直接返回None即可。

一切就绪以后，可以开始使用了，但是在这之前还需要确定你写的反向传播层的梯度否计算正确。pytorch提供了一个torch.autograd.gradcheck()函数来检查的所计算的梯度是否合理。这个检查的原理是通过比较梯度的数值计算和解析表达之间的误差来判断梯度计算是否正确：

梯度的数值计算法：

梯度的解析法就是我们通过求导公式计算得到的。如：

这一步检查无误以后就可以happy的使用这个模块了。也就是说完整的完成了一个pytorch的c++扩展。

总结一下：首先要写C++源码程序，需要使用一个torch的库<torch/extension.h>。这个库里面规定了如何利用这个库来写C++的extension，里面的基本数据格式为Tensor类型，不是一般的int/char/float类型。需要在C++源码中写一个forward函数和backward函数,在C++源码的最后使用PYBIND11来进行C++11和python的对话。之后便是编写setup.py文件，使用python提供的setuptools和torch自带的BuildExtension和CppExtension工具来进行编译的准备工作。然后再命令行中键入python setup.py install（根据需要使用build或者install，如你不想把你的package安装到系统路径中去，也就是site-package中，那么就用build命令，反之就用install命令），编译完成后还需要使用torch.autograd.Function来将这个扩展写成一个函数，方便在构建网络的时候调用。最后就在合适的地方使用Function.apply(*args)。这样一个定制化的模块就搞定了，也就是说完成了一个完整的pytorch的C++扩展（让用户丝毫感觉不到这个代码是在C++上扩展的~，但是作为一名算法菜狗，了解这些过程还是有必要的，这本身就是问题的一部分，包括如何写新网络中新的layer，毕竟看到再多，如果不自己动手始终差些感觉）

最后补充一点，关于求导，标量对标量的求导就不用多说了，这里主要是标量对向量/矩阵的求导。

这里使用到的标量对于矩阵的求导方法可以参见以下文章，写的非常好。

https://zhuanlan.zhihu.com/p/24709748

-END -

推荐阅读：

极市独家福利

40万奖金的AI移动应用大赛，参赛就有奖，入围还有额外奖励

添加极市小助手微信 （ID : cv-mart） ，备注： 研究方向-姓名-学校/公司-城市 （如：目标检测-小极-北大-深圳），即可申请加入 目标检测、目标跟踪、人脸、工业检测、医学影像、三维&SLAM、图像分割等极市技术交流群 ，更有 每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、 干货资讯汇总、行业技术交流 ，一起来让思想之光照的更远吧~