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

最新推荐文章于 2024-04-23 17:57:38 发布
最新推荐文章于 2024-04-23 17:57:38 发布
阅读量1.9k 收藏 13
点赞数 6
原文链接：https://mp.weixin.qq.com/s?__biz=Mzg2NzUxNTU1OA==&mid=2247519696&idx=2&sn=69c1d105ff0f9139f418487fb5a34edd&chksm=ceb8a019f9cf290fb6bd4a909cd91d517c1284558d896770092ad18355615408e261c1b25261&scene=126&sessionid=0
版权

作者 | 雅痞  编辑 | 汽车人

原文链接：https://zhuanlan.zhihu.com/p/595851188

点击下方卡片，关注“自动驾驶之心”公众号

ADAS巨卷干货，即可获取

点击进入→自动驾驶之心【模型部署】技术交流群

后台回复【模型部署工程】获取基于TensorRT的分类、检测任务的部署源码！

前段时间一直在做算子上的优化加速工作，在和其他同学的讨论中发现用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]。

参考

  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

06642d678e971d4c455b84815f0ed2ee.png

【自动驾驶之心】全栈技术交流群

自动驾驶之心是首个自动驾驶开发者社区，聚焦目标检测、语义分割、全景分割、实例分割、关键点检测、车道线、目标跟踪、3D目标检测、BEV感知、多传感器融合、SLAM、光流估计、深度估计、轨迹预测、高精地图、NeRF、规划控制、模型部署落地、自动驾驶仿真测试、硬件配置、AI求职交流等方向；

9d782179327334e43d243ea932b149c4.jpeg

添加汽车人助理微信邀请入群

备注：学校/公司+方向+昵称

自动驾驶之心
关注
  • 6
    点赞
  • 踩
  • 13
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
Pytorch 实现sobel算子的卷积操作详解
09-18
本文将深入讲解如何使用PyTorch实现Sobel算子的卷积操作，以进行边缘检测。 Sobel算子是一种一阶微分算子，常用于图像边缘检测，因为它可以有效地捕捉图像中的梯度变化。Sobel算子包含两个方向的核：水平和垂直，...
pytorch中使用cuda扩展的实现示例
09-17
本文将详细介绍如何在PyTorch中实现一个CUDA扩展，以实现一个element-wise的加法操作。 首先，我们需要创建CUDA编程的源文件和头文件。在`mathutil_cuda_kernel.cu`中，我们定义了两个关键部分：GPU执行函数`...
PyTorch自定义C++拓展算子、转换ONNX模型、ORT部署实现
嘤熊嚎杰的博客
03-01 1398
我们先来看一下建立的PyTorch模型，这是一个三层一维卷积模型，最后接了线性层完成分类任务输出。可以看到将应有的激活函数替换为了LIFSpike，并且传入了thresh和beta两个参数，模型中的其他参数省略。以下是LIFSpike的python部分实现，完整代码请见SNN-HAR，这里为了方便理解把源码中的ZIF替换为了FastSigmoid。
CUDA实战 | CUDA卷积算子手写详细实现流程！
CV_Autobot的博客
04-23 1497
作者|Pegessi 编辑| 极市平台原文链接：https://zhuanlan.zhihu.com/p/613538649点击下方卡片，关注“自动驾驶之心”公众号ADAS巨卷干货，即可获取点击进入→自动驾驶之心【模型部署】技术交流群后台回复【CUDA】获取CUDA实战书籍！导读本篇文章主要介绍如何利用CUDA实现一个2D卷积算子，实现过程较为简单，最终的实现效果可以在较小的尺寸下取得比c...
cudnn实现前向卷积运算（入门向）
qq_58135699的博客
04-23 771
注意有个小坑：cudnn的更新比较快所以在自己的环境下往往需要对算法进行修改，主要是同样功能的函数名称会有变化。可以去英伟达官网的cudnn网页查看各个API的功能。最后上面的程序运行结果是这样哒！算法的设计流程大致如下。
【CUDA编程--编程模型简介&算子开发流程】
xzpdxz的博客
11-13 2183
CUDA全称（Compute Unified Device Architecture）统一计算架构，是NVIDIA推出的并行计算平台深度学习加速：对于神经网络，无论是离线训练还是在线推理，都有巨量的矩阵、归一化、softmax等运算，且其中有非常多的并行计算，非常适合用GPU来进行运算加速一般来说，应用程序混合有并行部分和顺序部分，因此系统设计时混合使用 GPU 和 CPU，以最大限度地提高整体性能。具有高度并行性的应用程序可以利用 GPU 的大规模并行特性来实现比 CPU 更高的性能。
PyTorch 框架 AICPU 算子开发全流程
qq_44821958的博客
05-09 2721
PyTorch 框架 AICPU 算子开发全流程 文章目录PyTorch 框架 AICPU 算子开发全流程一．AICPU算子基本概念介绍二．算子开发流程介绍三. LogSpace算子分析四. LogSpace算子工程创建五．LogSpace算子原型定义和代码开发六．LogSpace算子逻辑实现部分代码开发七．LogSpace算子编译和部署八．LogSpace算子UT测试九．LogSpace算子ST测试十．经验总结十一．关于MindStudio更多的内容 详情指路视频专栏：https://www.bilibi
c++扩展算子开发③：CUDA算子的开发
m0_63642362的博客
01-22 1356
c++扩展算子开发③：CUDA算子的开发
Anchor-free目标检测论文汇总
3D视觉工坊
11-09 701
编辑丨CV技术指南前言 Anchor-free 目标检测是目标检测近几年的主流趋势之一，本文将分享一个汇总了最近几年所有Anchor-free论文的github项目。Anchor-free...
Pytorch 使用不同版本的cuda的方法步骤
09-17
例如，如果你的CUDAExtension需要使用CUDA 9.0进行编译，可以在`setup.py`文件中添加以下内容： ```python from torch.utils.cpp_extension import BuildExtension, CUDAExtension ext = CUDAExtension(name='my_...
pytorch-softdtw-cuda:使用Numba对PyTorch进行（可微分）软动态时间规整的快速CUDA实现
04-14
forward()和backward()传递都使用CUDA实现。 我的实现部分受到启发，其中提出了基于对角线的Belman递归实现。 入门 此代码取决于和 。 只需在您的项目中包含soft_dtw_cuda.py ，就可以了！ 您还可以运行随附的...
NVIDIA驱动、CUDA和Pytorch及其依赖
最新发布
05-18
说明： 1、离线环境下安装NVIDIA驱动、CUDA和Pytorch（HUAWEI Kunpeng 920 + NVIDIA A100 + Kylin V10 SP2） 2、因为文件太大，所以资源中只包括PyTorch及其依赖，其他文件评论我，我会私发你。
PyTorch中的CUDA操作
NLP工程化
08-24 4093
CUDA(Compute Unified Device Architecture)是NVIDIA推出的异构计算平台，PyTorch中有专门的模块torch.cuda来设置和运行CUDA相关操作。
在PyTorch中使用CUDA, pytorch与cuda不同版本对应安装指南，查看CUDA版本，安装对应版本pytorch
热门推荐
qq_35831906的博客
11-11 5万+
在PyTorch中使用CUDA，根据你的具体环境和需求调整版本号，确保安装的PyTorch版本与你的CUDA版本兼容。本文介绍PyTorch中使用CUDA的一般步骤和对应cuda版本的pytorch安装指南。
docker环境中实现可变卷积（deformable conv）遇到的问题
allen_hai的博客
07-22 566
解决docker环境中实现可变卷积（deformableconv）遇到的问题。
【机器学习】PyTorch、Cuda 的安装和示例代码
黎明总是如期而至
08-04 1万+
PyTorch是一个开源的Python深度学习框架，可以用于构建各种类型的神经网络模型。
CUDA_CHECK(cudaFree(...))报错CUDA error 1
qq_35037684的博客
08-02 1062
GPT-3.5太好用了，报错情况如下：
【代码小记】赏析《自定义C++和CUDA扩展》
qq_42889517的博客
05-17 313
参考文档：https://www.bookstack.cn/read/PyTorch-1.2-zh/3b1f3b4c60b8157d.md 【代码注释】 #编写CUDA扩展的一般策略： #①编写c++文件（correlation.cpp）：先定义将从Python调用的函数，在利用pybind11将C++函数或类衔接到Python中。 #②该c++文件（correlation.cpp）声明在CUDA（correlation_kernel.cu）文件中定义的函数。 #③c++函数将进行一些检查，并最终将其调.
pytorch模型量化
DeepLearning_的博客
11-06 2268
********************* range_trackers(范围统计器，统计量化前范围) *********************# ********************* bn融合_量化卷积（bn融合后，同时量化A/W，并做卷积） *********************即将量化后的值乘以就得到了反量化的结果，当然这个过程是有损的，如下图所示，橙色线表示的就是量化前的范围，而蓝色线代表量化后的数据范围，注意权重取。其中表示量化的缩放因子，和分别表示量化前和量化后的数值。
如何用cuda c写卷积算子
03-29
以下是一个简单的使用CUDA C编写卷积操作的示例代码： ```c #include <stdio.h> #include <stdlib.h> #include <cuda_runtime.h> #define BLOCK_SIZE 16 __global__ void convolve(float *input, float *output, float *kernel, int input_width, int input_height, int kernel_size) { int row = blockIdx.y * blockDim.y + threadIdx.y; int col = blockIdx.x * blockDim.x + threadIdx.x; float sum = 0.0f; for (int i = 0; i < kernel_size; ++i) { for (int j = 0; j < kernel_size; ++j) { int input_row = row + i - kernel_size / 2; int input_col = col + j - kernel_size / 2; if (input_row >= 0 && input_row < input_height && input_col >= 0 && input_col < input_width) { sum += input[input_row * input_width + input_col] * kernel[i * kernel_size + j]; } } } output[row * input_width + col] = sum; } int main() { int input_width = 512; int input_height = 512; int kernel_size = 5; float *input = (float*)malloc(input_width * input_height * sizeof(float)); for (int i = 0; i < input_width * input_height; ++i) { input[i] = rand() % 256 / 255.0f; } float *kernel = (float*)malloc(kernel_size * kernel_size * sizeof(float)); for (int i = 0; i < kernel_size * kernel_size; ++i) { kernel[i] = rand() % 256 / 255.0f; } float *output = (float*)malloc(input_width * input_height * sizeof(float)); float *d_input, *d_kernel, *d_output; cudaMalloc(&d_input, input_width * input_height * sizeof(float)); cudaMalloc(&d_kernel, kernel_size * kernel_size * sizeof(float)); cudaMalloc(&d_output, input_width * input_height * sizeof(float)); cudaMemcpy(d_input, input, input_width * input_height * sizeof(float), cudaMemcpyHostToDevice); cudaMemcpy(d_kernel, kernel, kernel_size * kernel_size * sizeof(float), cudaMemcpyHostToDevice); dim3 block_size(BLOCK_SIZE, BLOCK_SIZE); dim3 grid_size((input_width + BLOCK_SIZE - 1) / BLOCK_SIZE, (input_height + BLOCK_SIZE - 1) / BLOCK_SIZE); convolve<<<grid_size, block_size>>>(d_input, d_output, d_kernel, input_width, input_height, kernel_size); cudaMemcpy(output, d_output, input_width * input_height * sizeof(float), cudaMemcpyDeviceToHost); cudaFree(d_input); cudaFree(d_kernel); cudaFree(d_output); free(input); free(kernel); free(output); return 0; } ``` 该示例代码中，`convolve()`函数是卷积操作的核函数，它接收输入图像、输出图像和卷积核作为输入参数，以及输入图像尺寸和卷积核尺寸。在核函数中，每个线程负责计算输出图像中的一个像素值，它将卷积核与输入图像中对应像素的值相乘，并将结果累加到一个变量中。最后，输出图像中对应像素的值被赋为累加的结果。 在主函数中，我们首先生成随机的输入图像和卷积核，然后在GPU上分配内存，将输入图像和卷积核从主机内存复制到设备内存中，调用卷积核函数进行卷积操作，最后将输出图像从设备内存复制到主机内存中，并释放分配的内存。 在实际使用中，您需要根据自己的需求修改输入图像、卷积核和卷积核函数，以适应不同的场景。

“相关推荐”对你有帮助么？

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

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

抵扣说明：

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

余额充值