基于pytorch 1.1.0编写cuda扩展
0. 前言对于一些特殊的算子, 我们需要进行定制其前向和反向的过程, 从而使得其能够获得更快的速度, 加速模型的训练. 这样, 我们自然会想到使用PyTorch的cuda扩展来实现, 这里, 我将以一个简单且易于理解的例子出发, 详细的介绍如何构造一个属于你的cuda扩展.
1. 为什么需要写cuda扩展?
由于我们的一些特殊结构可以由基础的pytorch提供的算子进行组合而形成, 但是, 其问题是:
虽然已经使用了NVIDIA cuDNN、Intel MKL和NNPACK这些底层来加快训练速度，但是在某些情况下，比如我们要实现一些特定算法，光靠组合pytorch已有的操作是不够的。这是因为pytorch虽然在特定操作上经过了很好的优化，但是对于pytorch已经写好的这些操作，假如我们组合起来，组成我们的新的算法，pytorch才不管你的算法的具体执行流程，一般pytorch只会按照设计好的操作去使用GPU的通道，然后通道不能充分利用或者直接超负载，然后python解释器也不能对此进行优化，导致程序执行速度反而变慢了。
即由于一些自定义的操作是多个基础算子的组集, 这不可避免的导致吞吐量变小, 中间步骤变的繁多, 并没有充分利用硬件性能. 因此, 较好的方式是将复杂的操作进行fuse(也就是算子融合), 然后减少数据在多个操作之间流转,增强数据的本地性(locality), 从而提升GPU利用效率和计算速度.
这就是为什么我们需要写cuda扩展的原因: 对复杂逻辑的定制化加速处理
2. 写一个最基础的cuda扩展需要什么内容?
一个最简单的cuda扩展, 我们以传入一个4维的Tensor为例, 对其加N进行说明(注: 本例将在下文展开).
文件结构(the simplest cuda extension for pytorch 1.x):
– setup.py 安装文件
– cuda_ext/ cuda扩展所在文件夹
---- test_cuda.cpp cuda扩展声明&python binding
---- test_cuda_kernel.cu cuda扩展实际代码
当我们写好test_cuda.cpp, test_cuda_kernel.cu(cuda扩展的实际内容), 最后写好setup.py这个表示安装的文件, 即可进行安装并在pytorch1.1.0中使用你写的扩展啦~
3. 简单的例子: 对4维张量逐元素加N(N是自己指定的值)
3.1 环境说明
由于gcc版本和cuda的不兼容, 如果你使用的是gcc 6.x以上的版本, 那么务必需要将你的cuda和cudnn升级, 其中: cuda需要升级到10.0, cudnn也做对应升级. 如果你使用的cuda为9.0, 会出现错误.
本文使用的环境如下:
Ubuntu18.04
gcc 7.4.0
cuda 10.0
cudnn 7.6.4
torch == 1.1.0
pybind11
3.2 声明文件: cuda_ext/test_cuda.cpp
这里, 我们定义了① test1 ② test1_cuda 2个函数, 其中test1是调用test1_cuda的, 而test1_cuda就是我们即将在3.3中介绍的cuda实现的声明文件, 即cuda_ext/test_cuda.cpp在这里如同一个声明文件, 具体的定义在cuda_ext/test_cuda_kernel.cu中.
#include <torch/torch.h>
/*
define your own cuda extension.
This example is just add N to the original tensor.
*/
// CUDA forward declarations
at::Tensor test1_cuda(
at::Tensor image,
size_t N);
// C++ interface
#define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
at::Tensor test1(
at::Tensor image,
size_t N) {
// 类型检查. 必须要加.
CHECK_INPUT(image);
return test1_cuda(image, N);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("test1", &test1, "test1 (CUDA)");
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
需要注意的是, 我们在cuda_ext/test_cuda.cpp要把刚才写的test1函数绑定到module(这里的module在下面的setup.py定义, 为add_one_cuda)上，这里, 在cuda_ext/test_cuda.cpp文件最下面加上如下代码即可
// m是module的意思,不是method哦~, 这里的含义是, 为TORCH_EXTENSION_NAME模块绑定了名为test1的方法.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("test1", &test1, "test1 (CUDA)");
}
1
2
3
4
3.3 定义文件: cuda_ext/test_cuda_kernel.cu
这个部分是需要我们实现的具体的cuda编码内容, 由于我们的目的很简单: 对4维张量逐元素加N(N是自己指定的值), 所以其实现如下:
#include <ATen/ATen.h>
#include <cuda.h>
#include <cuda_runtime.h>
/*
define your own cuda extension.
This example is just add N to the original image.
*/
namespace {
template <typename scalar_t>
__global__ void test1_cuda_kernel(
scalar_t* __restrict__ image,
size_t N,
size_t batch_size,
size_t channel,
size_t image_height,
size_t image_width) {
int idx = blockDim.x * blockIdx.x + threadIdx.x;
int num_threads = blockDim.x * gridDim.x;
// 对每个element都加N, 看到image不是constant的, 我们直接对传入的4维张量做修改.
while(idx < batch_size*channel*image_height*image_width) {
image = image + N;
idx += num_threads;
}
}
}
at::Tensor test1_cuda(
at::Tensor image,
size_t N) {
const auto batch_size = image.size(0);
const auto channel = image.size(1);
const auto image_height = image.size(2);
const auto image_width = image.size(3);
const int threads = 32;
const dim3 blocks ((batch_size * channel - 1) / threads + 1);
// 注意, AT_DISPATCH_FLOATING_TYPES的第2个参数必须和所在函数体的名称一样! 否则会就无法dispatch.
AT_DISPATCH_FLOATING_TYPES(image.type(), "test1_cuda", ([&] {
test1_cuda_kernel<scalar_t><<<blocks, threads>>>(
image.data<scalar_t>(),
N,
batch_size,
channel,
image_height,
image_width);
}));
cudaError_t err = cudaGetLastError();
if (err != cudaSuccess)
printf("Error in test1: %s\n", cudaGetErrorString(err));
return image;
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
需要说明的点:
① template <typename scalar_t>的目的是为了泛化类型, 使得传入的pytorch tensor (cuda) 可以是单精度, 双精度 (int不行).
② cuda部分的核心在于并行, 这块不展开, 看test1_cuda_kernel的内容.
3.4 安装文件setup.py
好了, 现在, 我们把cuda_ext/下面的2个文件写好了, 那么我们需要按照第2部分所说的文件结构, 写一个setup.py, 其内容如下:
可以看出, 我们将写好的cuda扩展, 按照一定的方式, 用CUDAExtension进行封装, 然后扔进ext_modules里面作为setup函数的参数传入.
from setuptools import setup, find_packages
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
CUDA_FLAGS = []
ext_modules = [
# add_one_cuda 为 python调用的包名, 切记!
CUDAExtension('add_one_cuda', [
'cuda_ext/test1_cuda.cpp',
'cuda_ext/test1_cuda_kernel.cu',
])
]
INSTALL_REQUIREMENTS = ['numpy', 'torch', 'torchvision', 'scikit-image']
# https://pytorch.org/docs/master/cpp_extension.html
setup(
description='PyTorch implementation of <your own cuda extension>',
author='samuel ko', # 包的作者
author_email='samuel.gao023@gmail.com', # 包作者的邮箱
license='MIT License', # License类型
version='1.3.0', # 版本
name='add_one', # 包的名称
install_requires=INSTALL_REQUIREMENTS,# 预先需要的python依赖: numpy, torch等
ext_modules=ext_modules, # 这里指向我们的CUDAExtension.
cmdclass={'build_ext': BuildExtension}
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
写好之后, 我们就可以进行安装并验证了~
3.5 进行安装
安装的逻辑很简单, 进入你所想要的使用的虚拟环境, 然后执行
python3 setup.py install
1
即可, 但这里会出现一系列的问题, 我们把问题总结了一下,放在第4部分, 有需要的同学可以参考.
正确安装成功后, 会显示类似如下的信息:
3.6 进行验证
好了, 安装完之后, 我们可以验证自己写的这个扩展能否正常使用(这里有一点要强调的是: 引入的包的名字要和CUDAExtension中的名字一致, 而非setup函数里传的名字一致!):
# -*- coding: utf-8 -*-
import torch
import add_one_cuda
a = torch.randn(1, 3, 2, 2).cuda()
print(a)
# 调用我们写好的, 绑定到add_one_cuda模块上的cuda函数test1
print(add_one_cuda.test1(a, 5))
1
2
3
4
5
6
7
8
9
结果如下, 成功的对a逐元素的加5:
好了, 到这一步, 验证就全部通过了!
项目代码放在这里: github 地址, 大家可以直接拉下来跑一下.
4. 安装出现问题
4.1 error identifier __builtin_addressof is undefined
参考内容: https://blog.csdn.net/sophia_xw/article/details/100087061
解决方案:
sudo vim /usr/include/c++/7/bits/move.h (修改move.h)
在第44行开始添加如下内容:
4.2 RuntimeError: Ninja is required to load C++ extension
解决方案: 按顺序执行如下3条命令即可
wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip
sudo unzip ninja-linux.zip -d /usr/local/bin/
sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force
1
2
3
4.3 /usr/include/c++/6/tuple:495:244: error: wrong number of template arguments (4, should be 2)
原因: cuda版本过低 我当前的cuda版本是9.0,无法编译这套代码
解决: gcc版本与cuda9.0不兼容, 需要将cuda版本升级到10.0
cuda版本升级, 参考
参考资料
Pytorch拓展进阶(二)：Pytorch结合C++以及Cuda拓展
Error in cuda-extension compilation from pytorch advanced tutorial
Ubuntu下cuda版本升级
TORCH.UTILS.CPP_EXTENSION
Python包管理工具setuptools之setup函数参数详解
点赞
1
评论
分享
收藏
1
手机看
关注
收起全文
Keith
7716
pytorch学习笔记（十七）：python 端扩展 pytorch
pytorch 虽然提供了很多的 op 使得我们很容易的使用。但是当已有的 op 无法满足我们的要求的时候，那就需要自己动手来扩展。 pytorch 提供了两种方式来扩展 pytorch 的基础功能。 通过继承 autograd.Function通过 C 来扩展本篇博客主要介绍 继承 autograd.Function 来扩展 pytorch。官方文档链接继承 autograd.Functio...
tanmx219的博客
7273
pytorch通过torch.utils.cpp_extension构建CUDA/C++拓展
注意这个和前面的《Python与C语言混合编程：通过distutils或setuptools实现的一个简单的C扩展》不同，这个是pytorch的扩展，不是python的扩展。在pytorch的utils中，集成了setuptools模块。官方文档在这里：https://pytorch.org/docs/master/cpp_extension.html中文说明在这里：https://p......
...可/usr/include/c/6/tuple:495:244:error:wrong
6-10
编译nvcc报错‘/usr/include/c++/6/tuple’参考https://devtalk.nvidia.com/default/topic/1028112/cuda-setup-and-installation/nvcc-bug-related-to-gcc-6-lt-...
c++11中的tuple(元组)_C/C++_飞空静渡-CSDN博客
5-11
tuple看似简单,其实它是简约而不简单,可以说它是c++11中一个既简单又复杂的东东,关于它简单的一面是它很容易使用,复杂的一面是它内部隐藏了太多细节,要揭开它...
ROM的博客
9052
Pytorch学习笔记（一） 使用PyTorch搭建神经网络的套路
自古套路得人心！！！PyTorch Python中的张量和动态神经网络，具有强大的GPU加速功能。目录自古套路得人心！！！PyTorch目录常见模块的导入构建网络模型自定义Net类并实例化使用顺序容器构建网络并实例化定义训练函数定义测试函数常见模块的导入import torchimport torch.nn as nn......
g11d111的博客
2万+
PyTorch学习笔记(2)——变量类型(cpu/gpu)
前言PyTorch中的数据类型为Tensor，Tensor与Numpy中的ndarray类似，同样可以用于标量，向量，矩阵乃至更高维度上面的计算。PyTorch中的tensor又包括CPU上的数据类型和GPU上的数据类型，一般GPU上的Tensor是CPU上的Tensor加cuda()函数得到。通过使用Type函数可以查看变量类型。系统默认的torch.Tensor是torch.Floa......
解决thread:240: error: no matching function for cal..._CSDN博客
4-6
/usr/include/c++/7/thread:240: error: no matching function for call to ‘std::thread::_Invoker<std::tuple<void ()(), int> >::_M_invoke(std:...
C++11 std::tuple_c/c++_程序员Link-CSDN博客
5-24
C++中有 std::pair,C++11增加了std::tuple,是对std::pair的泛化。 一 简介 1 头文件<tuple> 2 声明 template<class ... Types> class tuple; 二tuple...
budding0828
356
pytorch 学习笔记4 —— GPU加速
在GPU上训练来源：https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html#training-an-image-classifier定义设备如果我们的机器支持CUDA，那么我们使用第一个cuda设备为cuda0device = torch.device("cuda:0" if torch.cuda.is_av...
FrankDura的博客
491
RuntimeError: Ninja is required to load C++ extension #167
报错RuntimeError: Ninja is required to load C++ extension #167解决wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zipsudo unzip ninja-linux.zip -d /usr/local/bin/sudo u...
dlhlSC的博客
1915
RuntimeError: Ninja is required to load C++ extension
wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zipsudo unzip ninja-linux.zip -d /usr/local/bin/sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin...
CodeTutor
2万+
PyTorch学习系列(十六)——如何使用cuda进行训练？
如果想在CUDA上进行计算，需要将操作对象放在GPU内存中。对于普通的张量，可以直接：x=torch.randn(2,3)x=x.cuda()对于神经网络：model=MyModel()model.cuda()同一个GPU上的张量计算后的结果仍然保存在该GPU上。参考http://pytorch.org/docs/notes/cuda.html#...
weixin_44207623的博客
5139
win 10+python3.6 安装NVIDIA apex
github下载源码https://github.com/NVIDIA/apex命令行：cd apexpython setup.py install [--cuda_ext] [--cpp_ext]遇到的问题：1.setuptools有ModuleNotFoundError→更新setuptoolspip install --upgrade setuptools2.[–cuda_...
jinbeibei0606的博客
203
ubuntu 16.04安装spatial-correlation-sampler出错error: command '/usr/local/cuda/bin/nvcc' failed with exi
记录问题环境：ubuntu16.04cuda-9.0pytorch 1.4gcc6.5安装spatial-correlation-sampler出现如下错误：/usr/include/c++/6/tuple: In instantiation of ‘static constexpr bool std::_TC<, _Elements>::_NonNestedTuple(...
qq_21407487的博客
402
CUDA9.0编译nvcc报错‘/usr/include/c++/6/tuple’
编译nvcc报错‘/usr/include/c++/6/tuple’参考https://devtalk.nvidia.com/default/topic/1028112/cuda-setup-and-installation/nvcc-bug-related-to-gcc-6-lt-tuple-gt-header-/可能原因GCC版本太高，更换低版本GCC编译器The latest g...
daniaokuye的专栏
546
pytorch0.4添加自定义cuda层
1. cuda函数的基础类型c或者c++pytorch扩展工具函数时，c或者c++是分开的。基本的命令和g++有很多相似之处。所以，打算使用不同的扩展方式时，需要相应的源码是c或者c++的。不然会出很多乱七八糟的问题。另外，pytorch 的官方网站上能找到完全用c语言写forward和backward的例子，我这里是将核心函数由cuda和c来编写，数据流的处理还是python下写的方式来。......
g11d111的博客
1235
PyTorch学习笔记(14) ——PyTorch 1.0 的C++ FrontEnd初体验
在去年12月份,我尝试了一下PyTorch 1.0的C++前端, 当时官方负责PyTorch的C++前端的老哥是: Peter Goldsborough, 当时的C++前端还不够稳定,官方文档提供的demo无法跑通.所以为了避免后面的同学再次入坑. 本部分将会手把手教你如何用PyTorch1.0 跑通C++ Mnist模型的训练~0. PyTorch C++接口哲学PyTorch’s ......
风海流的豪赌
520
PyTorch 学习笔记 (3) PyTorch 1.0+ C++/CUDA extension
参考教程PyTorch官方教程本机系统Ubuntu 18.04 LTSNVIDIA GeForce GTX 1080 with driver 430.50CUDA V10.1.243python 3.6.8 virtualenvtorch.version == ‘1.2.0’Visual Studio Code配置由于我使用的是python 的virtual environme...
h8832077的博客
2993
Windows 10 + Visual Studio 2017 + CUDA 10 环境下编译 pytorch 1.0
pytorch 1.0 一个月之前发布了。pytorch其实笔者很早就接触过，那时候惊叹于它的简洁、动态及良好的社区支持。但是那时候，pytorch在c++上的支持并不好，工业界很难用，基本上只属于一种比较好的算法验证框架。但是，pytorch 1.0 的发布终结了这一现况，现在，它有了良好的c++前向预测支持，再加上它本身不逊于tensorflow的运行速度，python式的开发逻辑（这点非常......
菜鸟求教：关于try catch(runtime_error RunTimeError)的问题
06-05
m0_37644085的博客
204
编译psroialign_cuda CHECK_CUDA(x); ‘CHECK_INPUT error: command 'gcc' failed with exit status 1
psroialign_cuda.cpp: In function ‘int psroialign_backward_cuda(at::Tensor, at::Tensor, at::Tensor, at::Tensor, float, int, int)’:psroialign_cuda.cpp:25:80: error: ‘AT_CHECK’ was not declared in this...
C++11可变参模板的一个问题,编译不过
06-07
Ninja Hacking- Unconventional Penetration Testing Tactics and Techniques 无水印pdf
09-25
c++编译出错
https://blog.csdn.net/g11d111/java/article/details/104407384
页:
[1]