像教女朋友一样教你用 Cuda 实现 PyTorch 算子

点击上方“小白学视觉”,选择加"星标"或“置顶

重磅干货,第一时间送达

前言

 

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类型要实现forwardbackward函数,并声明为静态成员函数。

forward

前向传播的前半部分就是正常传入Tensor进入接口,如果传入向量在之前的代码里是索引出来的很可能非连续,所以建议在传入算子的时候使其连续。

如果算子不需要考虑反向传播,可以用ctx.mark_non_differentiable(ans) 将函数的输出标记不需要微分[6]。

backward

backward的输入对应forward的输出输出对应forward的输入。例如这里backward的输入g_out对应forward输出ansbackward的输出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]。

参考

  1. CudaDemo代码仓库 https://github.com/Yuppie898988/CudaDemo

  2. PyTorch C++ API https://pytorch.org/cppdocs/api/library_root.html

  3. pybind11 https://pybind11.readthedocs.io/en/latest/basics.html

  4. setuptools https://setuptools.pypa.io/en/latest/userguide/quickstart.html#

  5. autograd https://pytorch.org/docs/stable/autograd.html

  6. mark_non_differentiable https://pytorch.org/docs/stable/generated/torch.autograd.function.FunctionCtx.mark_non_differentiable.html#torch.autograd.function.FunctionCtx.mark_non_differentiable

  7. save_for_backward https://pytorch.org/docs/stable/generated/torch.autograd.function.FunctionCtx.save_for_backward.html#torch.autograd.function.FunctionCtx.save_for_backward

  8. 详解PyTorch编译并调用自定义CUDA算子的三种方式 https://zhuanlan.zhihu.com/p/358778742

  9. PyTorch官方教程 https://pytorch.org/tutorials/advanced/cpp_extension.html

  10. 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“。请按照格式备注,否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告,否则会请出群,谢谢理解~
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值