Pytorch拓展进阶（二）：Pytorch结合C＋以及Cuda拓展

前言

之前的文章中：

Pytorch拓展进阶(一)：Pytorch结合C以及Cuda语言

。我们简单说明了如何简单利用C语言去拓展Pytorch并且利用编写底层的语言。这篇文章我们说明如何利用C++和Cuda去拓展Pytorch，同样实现我们的自定义功能。

为何使用C++

之前已经提到了什么我们要拓展，而不是直接使用Pytorch提供的python函数去构建算法函数。很简单因为效率和速度还有深度的自定义。

Pytorch虽然已经使用了NVIDIA cuDNN、Intel MKL和NNPACK这些底层来加快训练速度，但是在某些情况下，比如我们要实现一些特定算法，光靠组合Pytorch已有的操作是不够的。这是因为Pytorch虽然在特定操作上经过了很好的优化，但是对于Pytorch已经写好的这些操作，假如我们组合起来，组成我们的新的算法，Pytorch才不管你的算法的具体执行流程，一般Pytorch只会按照设计好的操作去使用GPU的通道，然后通道不能充分利用或者直接超负载，然后python解释器也不能对此进行优化，导致程序执行速度反而变慢了。

那么之前已经说到了利用c语言可以，那么C++拓展和C语言拓展的区别是什么呢？

C++是趋势。

即将出现，在官方的介绍中，有这么一段话：

TIM截图20180612144336

大致意思就是，C语言底层的库和C++底层的库会因为结合caffe2而有所改变，但是接口应该变动不会太大，上面提到了和比较耐人寻味。Aten是Pytorch现在使用的C++拓展专用库，Pytorch的设计者想去重构这个库以去适应caffe2.

那么，C++拓展的功能，相比C来说，应该是Pytorch更看重的一点(当然C还是能拓展的)，所以我们今天说一说C++拓展，长远来看，是值得去学习的。

以一个例子开始

同样，我们首先设计一个普通的神经网络层：

这个层叫做LLTM，即Long-Long-Term-Memory。一个经典的RNN构造。继承，然后按照我们平常的进行定义即可。

定义好了我们这样使用:

上面的程序当然是可以执行的，但是我们要注意，可以执行和效率是两码事，我们之所以要用拓展，是因为我们要提高我们算法运行的速度和实现我们所有想要实现自定义功能。

用C++写这个例子

在编写C++程序之前我们需要安装pybind这个python-C++拓展库。

.pybind11

安装pybind11很简单，执行下面两个命令就行：

注意：可能还需要安装和。

.用C++进行编写

好了，上面使我们的python版实现过程，现在我们改成C++版，当然我们首先编写一下简单的sigmoid功能函数：

上面的程序，引用的头文件很重要，这个头文件是在Pytorch安装后已经包含的，所以我们用尖括号括起来，这个头文件里面都有啥：

显然，这个头文件中大概有三个东西：

Aten库，这个库就是pytorch操作Tensor的C++接口

pybind11,这个是用来将python和C++结合起来

一些头文件，用来整合Aten和pybind11

好了，我们开始编写整个forward函数：

然后我们编写backward函数，需要注意的是Pytorch的C++接口并不会自动实现反向求导，需要我们自己去写，当然怎么写不用细究：

.绑定到Python

现在我们要把刚才写的C++绑定到Python上，我们在上面的文件最下面加上：

然后我们的文件目录大概是这样：

在中写这些信息，我们使用setuptools去编译我们的C++代码，和可以很方便地让我们实现对C++和Cuda的编译：

执行后就可以看到：

注意：编译有可能失败，这时候需要安装和等拓展组件。

.引用C++代码

编译好后，我们就可以：

然后我们就可以定义我们的自定义层了：

官方的结果中，程序在CPU和GPU中运行，C++版的运行速度要大于直接使用pytorch编写层的速度。注意，我们只编写了C++代码但是却可以在CPU中和GPU中跑，为什么，这就要归功于Aten的设计，Aten就是pytorch的C++版，使用Aten编写出来的tensor，只要在程序中，就可以将Tensor移到GPU当中了。

但是这样移到GPU中和直接编写cuda语言是不一样的。

编写CUDA代码

之前我们说明了如何写代码，现在我们来编写如何去写代码然后去和代码结合。C++和cuda代码结合其实和C语言是类似的，需要我们使用C++来写接口函数和python相连，然后使用C++去调用cuda程序。

首先我们写C++部分的接口程序：

注意上面的代码是接口程序，并没有具体的实现，而下面的代码才是核心的cuda代码：

只需要关注核心的代码即可，上面的代码后缀是，我们可以看到上面代码中也有C++的特性，那是因为cuda的nvcc编译器不仅支持c语言也支持C++语言的语法，我们需要注意的只是数据类型一定要写正确。返回的Tensor顺序不能错。在自己设计cuda程序的时候一定要注意内存问题。

具体的详细信心可以查看官方文档，这里只是简单介绍。

然后我们同样编写:

和之前不同的是我们采用的拓展变为了和包含文件包含了。

编译完成后，根据官网的性能测试，比起之前单纯使用C++在GPU上跑速度又提升了一阶。

官方代码地址:https://github.com/pytorch/extension-cpp

后记

使用C++和C都可以拓展pytorch实现自定义功能或者设计自己的算法。Pytorch拓展这些其实还是比较容易的，唯一的缺点就是，官方几乎没有这方面的文档说明，也就是接口说明，需要我们自己去研究。不过我们平时所需要的接口也就那么几种，多多编写熟悉了就就好。

参考资料

https://blog.csdn.net/fitzzhang/article/details/78988682

https://discuss.pytorch.org/t/compiling-an-extension-with-cuda-files/302/8

http://blog.christianperone.com/2018/03/pytorch-internal-architecture-tour/#python-intro

https://devblogs.nvidia.com/even-easier-introduction-cuda/

http://pybind11.readthedocs.io/en/master/index.html

https://pytorch.org/tutorials/advanced/cpp_extension.html

关注Oldpan博客，同步更新博客最新消息，持续酝酿深度学习质量文。