点击上方“小白学视觉”,选择加"星标"或“置顶”
重磅干货,第一时间送达
前言
CUDA(Compute Unified Device Architecture)是一种由NVIDIA推出的通用并行计算架构,该架构使GPU能够解决复杂的计算问题。开发人员可以使用C语言来为CUDA架构编写程序,所编写出的程序可以在支持CUDA的处理器上以超高性能运行。
编辑 | 自动驾驶之心 作者丨雅痞@知乎
链接丨https://zhuanlan.zhihu.com/p/595851188
前段时间一直在做算子上的优化加速工作,在和其他同学的讨论中发现用Cuda编写算子存在一定的门槛。虽然知乎上有很多优秀的教学指南、PyTorch官方也给出了tutorial(具体地址会放在文章末尾),但是对于每个环节的介绍与踩坑点似乎没有详实的说明。
结合我当时入门踩坑的惨痛经验,一个简单明了的demo能够大大减小上手的时间成本。所以我在这里以数组求和(下称sum_single
)、两数组相加(下称sum_double
) 为例,详细介绍一下用Cuda实现PyTorch算子的完整框架,具体的代码详见[1]:
https://github.com/Yuppie898988/CudaDemo
框架结构
├── ops
│ ├── __init__.py
│ ├── ops_py
│ │ ├── __init__.py
│ │ └── sum.py
│ └── src
│ ├── reduce_sum
│ │ ├── sum.cpp
│ │ └── sum_cuda.cu
│ └── sum_two_arrays
│ ├── two_sum.cpp
│ └── two_sum_cuda.cu
├── README.md
├── setup.py
└── test_ops.py
demo结构如上,其中
ops/src/
是Cuda/C++代码setup.py
是编译算子的配置文件ops/ops_py/
是用PyTorch包装的算子函数test_ops.py
是调用算子的测试文件
Cuda/C++
对于一个算子实现,需要用到.cu
(Cuda)编写核函数、.cpp
(C++)编写包装函数并调用PYBIND11_MODULE
对算子进行封装。
注意:Cuda文件和Cpp文件不能同名!!!否则编译不通过!!!
我们这里以src/sum_two_arrays/
为例进行解释
// src/sum_two_arrays/two_sum_cuda.cu
#include <cstdio>
#define THREADS_PER_BLOCK 256
#define WARP_SIZE 32
#define DIVUP(m, n) ((m + n - 1) / n)
__global__ void two_sum_kernel(const float* a, const float* b, float * c, int n){
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n){
c[idx] = a[idx] + b[idx];
}
}
void two_sum_launcher(const float* a, const float* b, float* c, int n){
dim3 blockSize(DIVUP(n, THREADS_PER_BLOCK));
dim3 threadSize(THREADS_PER_BLOCK);
two_sum_kernel<<<blockSize, threadSize>>>(a, b, c, n);
}
这里的关键是two_sum_kernel
这一核函数实现数组相加功能。下面的two_sum_launcher
函数负责分配线程块并调用核函数。
// src/sum_two_arrays/two_sum.cpp
#include <torch/extension.h>
#include <torch/serialize/tensor.h>
#define CHECK_CUDA(x) \
TORCH_CHECK(x.type().is_cuda(), #x, " must be a CUDAtensor ")
#define CHECK_CONTIGUOUS(x) \
TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous ")
#define CHECK_INPUT(x) \
CHECK_CUDA(x); \
CHECK_CONTIGUOUS(x)
void two_sum_launcher(const float* a, const float* b, float* c, int n);
void two_sum_gpu(at::Tensor a_tensor, at::Tensor b_tensor, at::Tensor c_tensor){
CHECK_INPUT(a_tensor);
CHECK_INPUT(b_tensor);
CHECK_INPUT(c_tensor);
const float* a = a_tensor.data_ptr<float>();
const float* b = b_tensor.data_ptr<float>();
float* c = c_tensor.data_ptr<float>();
int n = a_tensor.size(0);
two_sum_launcher(a, b, c, n);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &two_sum_gpu, "sum two arrays (CUDA)");
}
在C++文件中实现算子的封装,文件开头的宏定义函数是为了保证传入的向量在cuda上(CHECK_CUDA
)、传入的向量中元素地址连续(CHECK_CONTIGUOUS
)。two_sum_launcher
是对cuda文件中的声明。
two_sum_gpu
是与Python的接口,传入的参数是PyTorch中的Tensor。在这一部分需要对Tensor做CHECK检验(可选),并通过.data_ptr
得到Tensor变量的指针。对于Tensor在C++中的使用可查阅[2]。
最后PYBIND11_MODULE
的作用是对整个算子进行封装,能够通过Python调用C++函数[3]。对于自定义的其他算子,只用改动m.def
()中的三个参数
"forward"
:算子的方法名,假如算子的整个模块命名为sum_double
,则在Python中通过sum_double.forward
调用该算子&two_sum_gpu
:进行绑定的函数,这里根据自己实现的不同函数进行更改"sum two arrays (CUDA)"
:算子注释,在Python端调用help(sum_double.forward)
时会出现
可能有人会疑惑为什么要把算子和模块分开。假如整个
sum_double
有许多不同的功能,我就可以在一个模块中绑定多个算子,具体只用在PYBIND11_MODULE
中写入多个m.def()
,再通过sum_double.xxx
调用不同的算子
setup.py编译配置
在整个项目的根目录新建setup.py
文件配置编译信息,利用setuptools对算子打包[4]
from setuptools import find_packages, setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
name='CudaDemo',
packages=find_packages(),
version='0.1.0',
author='xxx',
ext_modules=[
CUDAExtension(
'sum_single', # operator name
['./ops/src/reduce_sum/sum.cpp',
'./ops/src/reduce_sum/sum_cuda.cu',]
),
CUDAExtension(
'sum_double',
['./ops/src/sum_two_arrays/two_sum.cpp',
'./ops/src/sum_two_arrays/two_sum_cuda.cu',]
),
],
cmdclass={
'build_ext': BuildExtension
}
)
文件中需要进行改动的有
name
:包名version
:包版本号author
:作者名称ext_modules
:编译C/C++扩展,list类型,每个元素为一个模块的相关信息(这里的模块在讲Cuda/C++这一块的末尾有提到,一个模块可以含有多个具体的算子)
CUDAExtension
在ext_modules
中采用CUDAExtension
指明Cuda/C++的文件路径 ,其中第一个参数为对应模块的名字,第二个参数为包含所有文件路径的列表。
这里的模块名和Cuda/C++中m.def()
定义的算子名共同决定了调用算子的方式。例如两数组相加的模块名是sum_double
、算子方法名是forward
, 所以在Python中调用该算子的方式为sum_double.forward()
。
值得一提的是packages
的值为list[str],表示本地需要打包的package。这里find_packages()
是找出本地所有的package。当然我们打包只用考虑ops/src/
中的文件,所以packages=['ops/src']
也能正常编译,不过为了方便还是采用find_packages()
。
PyTorch包装
为了让自定义算子能够正常正向传播、反向传播,我们需要继承torch.autograd.Function
进行算子包装[5]。我们这里以sum_double
为例进行介绍
# ops/ops_py/sum.py
import torch
from torch.autograd import Function
import sum_double
class SumDouble(Function):
@staticmethod
def forward(ctx, array1, array2):
"""sum_double function forward.
Args:
array1 (torch.Tensor): [n,]
array2 (torch.Tensor): [n,]
Returns:
ans (torch.Tensor): [n,]
"""
array1 = array1.float()
array2 = array2.float()
ans = array1.new_zeros(array1.shape)
sum_double.forward(array1.contiguous(), array2.contiguous(), ans)
# ctx.mark_non_differentiable(ans) # if the function is no need for backpropogation
return ans
@staticmethod
def backward(ctx, g_out):
# return None, None # if the function is no need for backpropogation
g_in1 = g_out.clone()
g_in2 = g_out.clone()
return g_in1, g_in2
sum_double_op = SumDouble.apply
文件开头import sum_double
就是导入的setup.py
中定义的模块名。
自定义的torch.autograd.Function
类型要实现forward
、backward
函数,并声明为静态成员函数。
forward
前向传播的前半部分就是正常传入Tensor进入接口,如果传入向量在之前的代码里是索引出来的很可能非连续,所以建议在传入算子的时候使其连续。
如果算子不需要考虑反向传播,可以用ctx.mark_non_differentiable(ans)
将函数的输出标记不需要微分[6]。
backward
backward
的输入对应forward
的输出,输出对应forward
的输入。例如这里backward
的输入g_out
对应forward
输出ans
,backward
的输出g_in1, g_in2
对应forward
输入array1, array2
。
如果算子不需要考虑反向传播,则直接return None, None
。否则就按照对应输入变量的梯度进行计算。
值得注意的是,如果反向传播需要用到forward
的信息,可以用ctx
进行记录存取。例如对一个数组求和,则反向传播的梯度为原数组长度的向量。就可以在forward
中用ctx.shape=array.shape[0]
记录输入数组长度,并在backward
中通过n=ctx.shape
进行读取。
如果存取的是Tensor则建议使用save_for_backward(x, y, z, ...)
存储向量,并用x, y, z, ...=ctx.saved_tensors
取向量,而不是直接用ctx[7]
。
To prevent incorrect gradients and memory leaks, and enable the application of saved tensor hooks.
注:save_for_backward()只能存向量,标量用ctx直接存取。
最后用sum_double_op = SumDouble.apply
获取最终的函数形式。
_init_.py
为了在外部调用包装好的PyTorch函数,通过ops/ops_py/__init__.py
声明
from .sum import sum_single_op, sum_double_op
__all__ = ['sum_single_op', 'sum_double_op']
ops/__init__.py
中
from .ops_py import *
Build & Test
提前安装好PyTorch环境,并在demo的根目录下pip install -e .
通过python test_ops.py
测试结果,没问题的情况应输出:
Average time cost of sum_single is 2.8257 ms
Average time cost of sum_double is 0.1128 ms
如果无法编译可能是没有将nvcc加入环境变量,ls /usr/local/
看看是否有cuda文件夹。例如我这里是cuda-11.6
文件夹,则进入~/.bashrc
在文件末尾加入
export PATH=$PATH:/usr/local/cuda-11.6/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-11.6/lib64
如果是没有配置cuda环境可以参考我之前的文章
https://zhuanlan.zhihu.com/p/558605762
其他相关参考资料可见[8][9][10]。
参考
CudaDemo代码仓库 https://github.com/Yuppie898988/CudaDemo
PyTorch C++ API https://pytorch.org/cppdocs/api/library_root.html
pybind11 https://pybind11.readthedocs.io/en/latest/basics.html
setuptools https://setuptools.pypa.io/en/latest/userguide/quickstart.html#
autograd https://pytorch.org/docs/stable/autograd.html
mark_non_differentiable https://pytorch.org/docs/stable/generated/torch.autograd.function.FunctionCtx.mark_non_differentiable.html#torch.autograd.function.FunctionCtx.mark_non_differentiable
save_for_backward https://pytorch.org/docs/stable/generated/torch.autograd.function.FunctionCtx.save_for_backward.html#torch.autograd.function.FunctionCtx.save_for_backward
详解PyTorch编译并调用自定义CUDA算子的三种方式 https://zhuanlan.zhihu.com/p/358778742
PyTorch官方教程 https://pytorch.org/tutorials/advanced/cpp_extension.html
THE C++ FRONTEND https://pytorch.org/cppdocs/frontend.html
下载1:OpenCV-Contrib扩展模块中文版教程
在「小白学视觉」公众号后台回复:扩展模块中文教程,即可下载全网第一份OpenCV扩展模块教程中文版,涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。
下载2:Python视觉实战项目52讲
在「小白学视觉」公众号后台回复:Python视觉实战项目,即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目,助力快速学校计算机视觉。
下载3:OpenCV实战项目20讲
在「小白学视觉」公众号后台回复:OpenCV实战项目20讲,即可下载含有20个基于OpenCV实现20个实战项目,实现OpenCV学习进阶。
交流群
欢迎加入公众号读者群一起和同行交流,目前有SLAM、三维视觉、传感器、自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN、算法竞赛等微信群(以后会逐渐细分),请扫描下面微信号加群,备注:”昵称+学校/公司+研究方向“,例如:”张三 + 上海交大 + 视觉SLAM“。请按照格式备注,否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告,否则会请出群,谢谢理解~