pytorch之C++实现自定义算子

自定义算子

对于输入 x ，其输出为

利用C++实现以上算子，总共只要实现两个文件：

setup.py

利用python中提供的setuptools模块完成事先编译流程，将写有算子的C++文件，编译成为一个动态链接库（在Linux平台是一个.so后缀文件），可以让python调用其中实现的函数功能。需要setup.py编写如下

from setuptools import setup
from torch.utils import cpp_extension

setup(
    name='ncrelu_cpp',						# 编译后的链接库名称
    ext_modules=[
        cpp_extension.CppExtension(
            'ncrelu_cpp', ['ncrelu.cpp']		       # 待编译文件，及编译函数
        )
    ],
    cmdclass={						       # 执行编译命令设置
        'build_ext': cpp_extension.BuildExtension
    }
)

这里Pytorch提供了一个封装cpp_extension，方便编译过程中所需要的设置选项，以及所需包含的头文件位置路径设置等等。

ncrelu.cpp

接下来便是C++实现的函数ncrelu.cpp。首先上代码

#include <torch/extension.h>					// 头文件引用部分 

torch::Tensor ncrelu_forward(torch::Tensor input) {
    auto pos = input.clamp_min(0);				       // 具体实现部分
    auto neg = input.clamp_max(0);
    return torch::cat({pos, neg}, 1);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {	// 绑定部分
    m.def("forward", &ncrelu_forward, "NCReLU forward");
}

以上代码包含了三个部分。分别是

头文件引用部分：这里包含了torch/extension.h头文件，是编写Pytorch的C++扩展时必须包含的一个文件。它基本上囊括了实现中所需要的所有依赖，包含了ATen库，pybind11和二者之间的交互。其中ATen是Pytorch底层张量运算库，负责实现具体张量操作运算；pybind11是实现C++代码到python的绑定（binding），可以在python里调用C++函数。

具体实现部分：函数返回类型和传递参数类型均是torch::Tensor类，这种对象不仅包含了数据，还附属了诸多运算操作。因此我们可以看到在下面实现方式类似于python中使用Pytorch张量运算操作一样，可以直接调用如截取操作clamp和拼接操作cat等，非常简洁已读且方便。

绑定部分：只需要在m.def中传入参数，分别是绑定到python的函数名称，需绑定的C++函数引用，以及一个简短的函数说明字符串，用来添加到python函数中的__doc__成员名称中。

将以上两个文件放在同一文件夹下，然后进行编译。使用python setup.py build_ext --inplace命令，如果一切正常将会在文件夹下产生类似ncrelu_cpp.cpython-35m-x86_64-linux-gnu.so动态链接文件。然后我们可以启动python检测是否可以导入其中的函数

>> import torch
>> import ncrelu_cpp
>> a = torch.randn(4, 3)
>> a
tensor([[ 0.5921,  0.3207,  0.7690],
        [ 1.4514, -0.8942,  0.9039],
        [-0.3262, -0.1610,  0.6137],
        [ 0.7824, -1.8527,  0.2844]])
>> b = ncrelu_cpp.forward(a)
>> b
tensor([[ 0.5921,  0.3207,  0.7690,  0.0000,  0.0000,  0.0000],
        [ 1.4514,  0.0000,  0.9039,  0.0000, -0.8942,  0.0000],
        [ 0.0000,  0.0000,  0.6137, -0.3262, -0.1610,  0.0000],
        [ 0.7824,  0.0000,  0.2844,  0.0000, -1.8527,  0.0000]])