像教女朋友一样教你用Cuda实现PyTorch算子 |
您所在的位置：网站首页 › pytorch的基本使用 › 像教女朋友一样教你用Cuda实现PyTorch算子 |
|
前段时间一直在做算子上的优化加速工作，在和其他同学的讨论中发现用Cuda编写算子存在一定的门槛。虽然知乎上有很多优秀的教学指南、PyTorch官方也给出了tutorial（具体地址会放在文章末尾），但是对于每个环节的介绍与踩坑点似乎没有详实的说明。 结合我当时入门踩坑的惨痛经验，一个简单明了的demo能够大大减小上手的时间成本。所以我在这里以数组求和（下称sum_single）、两数组相加（下称sum_double） 为例，详细介绍一下用Cuda实现PyTorch算子的完整框架，具体的代码详见[1] 框架结构├── 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.pydemo结构如上，其中 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 #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(a, b, c, n); }这里的关键是two_sum_kernel这一核函数实现数组相加功能。下面的two_sum_launcher函数负责分配线程块并调用核函数。 // src/sum_two_arrays/two_sum.cpp #include #include #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(); const float* b = b_tensor.data_ptr(); float* c = c_tensor.data_ptr(); 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]。 backwardbackward的输入对应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环境可以参考我之前的文章 希望这篇文章能帮到每一位想在PyTorch加入自定义Cuda算子的朋友，欢迎知乎点赞、讨论，或者去仓库点个star！ 其他相关参考资料可见[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 |
今日新闻 |
推荐新闻 |
| CopyRight 2018-2019 办公设备维修网 版权所有 豫ICP备15022753号-3 |