PyTorch中的C++扩展实现


Posted in Python onApril 02, 2020

今天要聊聊用 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 提供的<TH/TH.h><THC/THC.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 <torch/extension.h>
#include <vector>

// 前向传播
torch::Tensor Test_forward_cpu(const torch::Tensor& inputA,
              const torch::Tensor& inputB);
// 反向传播
std::vector<torch::Tensor> Test_backward_cpu(const torch::Tensor& gradOutput);

注意,这里引用的<torch/extension.h>头文件至关重要,它主要包括三个重要模块:

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

源文件 test.cpp 如下:

#include "test.h"

// 前向传播,两个 Tensor 相加。这里只关注 C++ 扩展的流程,具体实现不深入探讨。
torch::Tensor Test_forward_cpu(const torch::Tensor& x,
              const torch::Tensor& y) {
  AT_ASSERTM(x.sizes() == y.sizes(), "x must be the same size as y");
  torch::Tensor z = torch::zeros(x.sizes());
  z = 2 * x + y;
  return z;
}

// 反向传播
// 在这个例子中,z对x的导数是2,z对y的导数是1。
// 至于这个backward函数的接口(参数,返回值)为何要这样设计,后面会讲。
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};
}

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

第二步

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

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

以下是 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_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++ 代码:

python setup.py install

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

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

from torch.autograd import Function

import test_cpp

class TestFunction(Function):

  @staticmethod
  def forward(ctx, x, y):
    return test_cpp.forward(x, y)

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

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

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

class Function(with_metaclass(FunctionMeta, _C._FunctionBase, _ContextMethodMixin, _HookMixin)):
 
  ...

  @staticmethod
  def forward(ctx, *args, **kwargs):
    r"""Performs the operation.

    This function is to be overridden by all subclasses.

    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.
    """
    raise NotImplementedError

  @staticmethod
  def backward(ctx, *grad_outputs):
    r"""Defines a formula for differentiating the operation.

    This function is to be overridden by all subclasses.

    It must accept a context :attr:`ctx` as the first argument, followed by
    as many outputs did :func:`forward` return, and it should return as many
    tensors, as there were inputs to :func:`forward`. Each argument is the
    gradient w.r.t the given output, and each returned value should be the
    gradient w.r.t. the corresponding input.

    The context can be used to retrieve tensors saved during the forward
    pass. It also has an attribute :attr:`ctx.needs_input_grad` as a tuple
    of booleans representing whether each input needs gradient. E.g.,
    :func:`backward` will have ``ctx.needs_input_grad[0] = True`` if the
    first input to :func:`forward` needs gradient computated w.r.t. the
    output.
    """
    raise NotImplementedError

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

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

import torch

class Test(torch.nn.Module):

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

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

现在,我们的文件目录变成:

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

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

测试

下面,我们测试一下前向传播和反向传播:

import torch
from torch.autograd import Variable

from test import Test

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=<TestFunctionBackward>)
x.grad:  tensor([2., 2., 2.])
y.grad:  tensor([1., 1., 1.])

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

CUDA扩展

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

参考

CUSTOM C EXTENSIONS FOR PYTORCH
CUSTOM C++ AND CUDA EXTENSIONS
Pytorch拓展进阶(一):Pytorch结合C以及Cuda语言
Pytorch拓展进阶(二):Pytorch结合C++以及Cuda拓展

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

Python 相关文章推荐
Python 面向对象 成员的访问约束
Dec 23 Python
Python实现的多线程端口扫描工具分享
Jan 21 Python
python 使用get_argument获取url query参数
Apr 28 Python
Django项目实战之用户头像上传与访问的示例
Apr 21 Python
python 调用钉钉机器人的方法
Feb 20 Python
导入tensorflow时报错:cannot import name 'abs'的解决
Oct 10 Python
python 子类调用父类的构造函数实例
Mar 12 Python
Python使用sqlite3模块内置数据库
May 07 Python
Python实现在线批量美颜功能过程解析
Jun 10 Python
python判断all函数输出结果是否为true的方法
Dec 03 Python
python gui开发——制作抖音无水印视频下载工具(附源码)
Feb 07 Python
详解Python flask的前后端交互
Mar 31 Python
python实现将列表中各个值快速赋值给多个变量
Apr 02 #Python
Python运行提示缺少模块问题解决方案
Apr 02 #Python
Pycharm配置PyQt5环境的教程
Apr 02 #Python
Python无头爬虫下载文件的实现
Apr 02 #Python
linux 下selenium chrome使用详解
Apr 02 #Python
Python HTTP下载文件并显示下载进度条功能的实现
Apr 02 #Python
python实现将range()函数生成的数字存储在一个列表中
Apr 02 #Python
You might like
PHP入门教程之字符串处理技巧总结(转换,过滤,解析,查找,截取,替换等)
2016/09/11 PHP
php 删除一维数组中某一个值元素的操作方法
2018/02/01 PHP
thinkphp3.2同时连接两个数据库的简单方法
2019/08/13 PHP
PHPstorm激活码2020年5月13日亲测有效
2020/09/17 PHP
Jquery iframe内部出滚动条
2010/02/11 Javascript
鼠标左键单击冲突的问题解决方法(防止冒泡)
2014/05/14 Javascript
JS使用正则实现去掉字符串左右空格的方法
2016/12/27 Javascript
使用jquery给新生的th绑定hover事件的实例
2017/02/10 Javascript
Node学习记录之cluster模块
2017/05/31 Javascript
Vue.js常用指令的使用小结
2017/06/23 Javascript
基于JS脚本语言的基础语法详解
2017/07/22 Javascript
JavaScript递归算法生成树形菜单
2017/08/15 Javascript
在 webpack 中使用 ECharts的实例详解
2018/02/05 Javascript
vue中mint-ui的使用方法
2018/04/04 Javascript
Vue自定义表单内容检查rules实例
2020/10/30 Javascript
[02:40]DOTA2英雄基础教程 炼金术士
2013/12/23 DOTA
python用10行代码实现对黄色图片的检测功能
2015/08/10 Python
详解python之配置日志的几种方式
2017/05/22 Python
Python使用当前时间、随机数产生一个唯一数字的方法
2017/09/18 Python
TensorFlow实现AutoEncoder自编码器
2018/03/09 Python
用python生成(动态彩色)二维码的方法(使用myqr库实现)
2019/06/24 Python
Django 实现admin后台显示图片缩略图的例子
2019/07/28 Python
python反转列表的三种方式解析
2019/11/08 Python
tensorflow之获取tensor的shape作为max_pool的ksize实例
2020/01/04 Python
利用HTML5 Canvas制作键盘及鼠标动画的实例分享
2016/03/15 HTML / CSS
网络艺术零售业的先驱者:artrepublic
2017/09/26 全球购物
独特的礼品和创新的科技产品:The Grommet
2018/02/24 全球购物
精美的手工家居和生活用品:Nkuku
2019/11/01 全球购物
分层教学实施方案
2014/03/19 职场文书
婚纱店策划方案
2014/05/22 职场文书
小学先进集体事迹材料
2014/05/31 职场文书
人口与计划生育目标管理责任书
2014/07/29 职场文书
教师党员自我评议不足范文
2014/10/19 职场文书
行政司机岗位职责
2015/04/10 职场文书
PHP新手指南
2021/04/01 PHP
利用JuiceFS使MySQL 备份验证性能提升 10 倍
2022/03/17 MySQL