PyTorch中的C++扩展实现

2023-07-27 18:25:14 80

今天要聊聊用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并非难事。而随着PyTorch1.0的发布，官方已经开始考虑将PyTorch的底层代码用caffe2替换，因此他们也在逐步重构ATen，后者是目前PyTorch使用的C++扩展库。总的来说，C++是未来的趋势。至于CUDA，这是几乎所有深度学习系统在构建之初就采用的工具，因此CUDA的扩展接口是标配。

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

PyTorch的C、C++、CUDA扩展

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

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

C++扩展

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

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

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

第一步

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

头文件test.h

#include
#include

//前向传播
torch::TensorTest_forward_cpu(consttorch::Tensor&inputA,
consttorch::Tensor&inputB);
//反向传播
std::vectorTest_backward_cpu(consttorch::Tensor&gradOutput);

注意，这里引用的头文件至关重要，它主要包括三个重要模块：

pybind11，用于C++和python交互；
ATen，包含Tensor等重要的函数和类；
一些辅助的头文件，用于实现ATen和pybind11之间的交互。

源文件test.cpp如下：

#include"test.h"

//前向传播，两个Tensor相加。这里只关注C++扩展的流程，具体实现不深入探讨。
torch::TensorTest_forward_cpu(consttorch::Tensor&x,
consttorch::Tensor&y){
AT_ASSERTM(x.sizes()==y.sizes(),"xmustbethesamesizeasy");
torch::Tensorz=torch::zeros(x.sizes());
z=2*x+y;
returnz;
}

//反向传播
//在这个例子中，z对x的导数是2，z对y的导数是1。
//至于这个backward函数的接口（参数，返回值）为何要这样设计，后面会讲。
std::vectorTest_backward_cpu(consttorch::Tensor&gradOutput){
torch::TensorgradOutputX=2*gradOutput*torch::ones(gradOutput.sizes());
torch::TensorgradOutputY=gradOutput*torch::ones(gradOutput.sizes());
return{gradOutputX,gradOutputY};
}

//pybind11绑定
PYBIND11_MODULE(TORCH_EXTENSION_NAME,m){
m.def("forward",&Test_forward_cpu,"TESTforward");
m.def("backward",&Test_backward_cpu,"TESTbackward");
}

第二步

新建一个编译安装的配置文件setup.py，文件目录安排如下：

└──csrc
├──cpu
│├──test.cpp
│└──test.h
└──setup.py

以下是setup.py中的内容：

fromsetuptoolsimportsetup
importos
importglob
fromtorch.utils.cpp_extensionimportBuildExtension,CppExtension

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

setup(
name='test_cpp',#模块名称，需要在python中调用
version="0.1",
ext_modules=[
CppExtension('test_cpp',sources=source_cpu,include_dirs=[include_dirs]),
],
cmdclass={
'build_ext':BuildExtension
}
)

注意，这个C++扩展被命名为test_cpp，意思是说，在python中可以通过test_cpp模块来调用C++函数。

第三步

在cpu这个目录下，执行下面的命令编译安装C++代码：

pythonsetup.pyinstall

之后，可以看到一堆输出，该C++模块会被安装在python的site-packages中。

完成上面几步后，就可以在python中调用C++代码了。在PyTorch中，按照惯例需要先把C++中的前向传播和反向传播封装成一个函数op（以下代码放在test.py文件中）：

fromtorch.autogradimportFunction

importtest_cpp

classTestFunction(Function):

@staticmethod
defforward(ctx,x,y):
returntest_cpp.forward(x,y)

@staticmethod
defbackward(ctx,gradOutput):
gradX,gradY=test_cpp.backward(gradOutput)
returngradX,gradY

这样一来，我们相当于把C++扩展的函数嵌入到PyTorch自己的框架内。

我查看了这个Function类的代码，发现是个挺有意思的东西：

classFunction(with_metaclass(FunctionMeta,_C._FunctionBase,_ContextMethodMixin,_HookMixin)):

...

@staticmethod
defforward(ctx,*args,**kwargs):
r"""Performstheoperation.

Thisfunctionistobeoverriddenbyallsubclasses.

Itmustacceptacontextctxasthefirstargument,followedbyany
numberofarguments(tensorsorothertypes).

Thecontextcanbeusedtostoretensorsthatcanbethenretrieved
duringthebackwardpass.
"""
raiseNotImplementedError

@staticmethod
defbackward(ctx,*grad_outputs):
r"""Definesaformulafordifferentiatingtheoperation.

Thisfunctionistobeoverriddenbyallsubclasses.

Itmustacceptacontext:attr:`ctx`asthefirstargument,followedby
asmanyoutputsdid:func:`forward`return,anditshouldreturnasmany
tensors,astherewereinputsto:func:`forward`.Eachargumentisthe
gradientw.r.tthegivenoutput,andeachreturnedvalueshouldbethe
gradientw.r.t.thecorrespondinginput.

Thecontextcanbeusedtoretrievetensorssavedduringtheforward
pass.Italsohasanattribute:attr:`ctx.needs_input_grad`asatuple
ofbooleansrepresentingwhethereachinputneedsgradient.E.g.,
:func:`backward`willhave``ctx.needs_input_grad[0]=True``ifthe
firstinputto:func:`forward`needsgradientcomputatedw.r.t.the
output.
"""
raiseNotImplementedError

这里需要注意一下backward的实现规则。该接口包含两个参数：ctx是一个辅助的环境变量，grad_outputs则是来自前一层网络的梯度列表，而且这个梯度列表的数量与forward函数返回的参数数量相同，这也符合链式法则的原理，因为链式法则就需要把前一层中所有相关的梯度与当前层进行相乘或相加。同时，backward需要返回forward中每个输入参数的梯度，如果forward中包括n个参数，就需要一一返回n个梯度。所以，在上面这个例子中，我们的backward函数接收一个参数作为输入（forward只输出一个变量），并返回两个梯度（forward接收上一层两个输入变量）。

定义完Function后，就可以在Module中使用这个自定义op了：

importtorch

classTest(torch.nn.Module):

def__init__(self):
super(Test,self).__init__()

defforward(self,inputA,inputB):
returnTestFunction.apply(inputA,inputB)

现在，我们的文件目录变成：

├──csrc
│├──cpu
││├──test.cpp
││└──test.h
│└──setup.py
└──test.py

之后，我们就可以将test.py当作一般的PyTorch模块进行调用了。

测试

下面，我们测试一下前向传播和反向传播：

importtorch
fromtorch.autogradimportVariable

fromtestimportTest

x=Variable(torch.Tensor([1,2,3]),requires_grad=True)
y=Variable(torch.Tensor([4,5,6]),requires_grad=True)
test=Test()
z=test(x,y)
z.sum().backward()
print('x:',x)
print('y:',y)
print('z:',z)
print('x.grad:',x.grad)
print('y.grad:',y.grad)

输出如下：

x: tensor([1.,2.,3.],requires_grad=True)
y: tensor([4.,5.,6.],requires_grad=True)
z: tensor([6., 9.,12.],grad_fn=)
x.grad: tensor([2.,2.,2.])
y.grad: tensor([1.,1.,1.])

可以看出，前向传播满足z=2x+y，而反向传播的结果也在意料之中。

CUDA扩展

虽然C++写的代码可以直接跑在GPU上，但它的性能还是比不上直接用CUDA编写的代码，毕竟ATen没法并不知道如何去优化算法的性能。不过，由于我对CUDA仍一窍不通，因此这一步只能暂时略过，留待之后补充～囧～。

参考

CUSTOMCEXTENSIONSFORPYTORCH
CUSTOMC++ANDCUDAEXTENSIONS
Pytorch拓展进阶(一)：Pytorch结合C以及Cuda语言
Pytorch拓展进阶(二)：Pytorch结合C++以及Cuda拓展

到此这篇关于PyTorch中的C++扩展实现的文章就介绍到这了,更多相关PyTorchC++扩展内容请搜索毛票票以前的文章或继续浏览下面的相关文章希望大家以后多多支持毛票票！

PyTorch中的C++扩展实现

热门推荐

随机推荐