Pytorch-c++-CUDA

一、CUDA1.1 简介CUDACompute Unified Device Architecture是 NVIDIA 推出的并行计算平台和编程模型它让开发者能够利用 GPU 中成百上千个核心进行通用计算而不仅仅是图形渲染大幅加速科学计算、深度学习、图像处理等数据密集型任务。1.2 cuda的优势高并行吞吐量数千个线程同时运行适合处理大规模数据并行问题。细粒度的内存层次可利用共享内存、寄存器等实现极高的访存带宽。与 C/C 紧密集成开发者可快速上手将计算密集型部分卸载到 GPU 执行显著缩短运算时间1.3 自定义cuda的场景非平行计算(如体积渲染)大量串列计算。当应用的计算模式无法被现有库或简单并行化覆盖时就需要编写自定义 CUDA 核函数。典型场景如体积渲染虽然每条光线内部包含大量串行步进和采样计算但不同像素、不同光线之间却高度独立。通过自定义 CUDA 核函数可以将这种“每个元素内部串行、元素之间并行”的任务在 GPU 上高效执行将数以千计的串列计算过程同时推进从而获得远高于纯 CPU 串行计算的性能。二打通Pyton调用C文件2.1 使用VSCode配置环境桥梁C环境新建C文件ctrl shift p 弹出选项卡点击Edit Configurations UI配置如下图所示Edit Configuration点击后退出此时会发现文件夹中多了一个.vscode文件夹里面有c_cpp_properties.json文件打开后对includePath进行编辑根据自己的创建的python环境来定可以在你环境文件夹的目录下问gpt让他帮你找。我的环境名是GSRaw。{ configurations: [ { name: Win32, // C项目中编译时头文件搜索目录 includePath: [ ${workspaceFolder}/**, // 提供 CUDA 运行时 API C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.8/include, // py文件环境下的通用 Python/C 扩展头文件确保能与 Python 解释器交互 D:/biancheng/Python/Anaconda3_envs_dirs/GSRaw/include, // PyTorch C 核心库LibTorch头文件 D:/biancheng/Python/Anaconda3_envs_dirs/GSRaw/Lib/site-packages/torch/include, // PyTorch C 前端 API提供高层模块化接口方便构建训练/推理模型 D:/biancheng/Python/Anaconda3_envs_dirs/GSRaw/Lib/site-packages/torch/include/torch/csrc/api/include ], // 相当于字典用来排除单词下面波浪线的可不管 defines: [ _DEBUG, UNICODE, _UNICODE, __CUDACC__ ] } ], version: 4 }2.2 编写桥梁的C文件这里只是示例 用于python调用。文件名为interpolation.cpp# include torch/extension.h // 根据个人使用cuda的需求自定义的函数 torch::Tensor trilinear_interpolation( torch::Tensor feats, torch::Tensor point ){ return feats; } // 固定结构搭配 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){ // 别名setup.py文件加入环境后py文件的可以引用的包名 上面自定义的函数名 m.def(trilinear_interpolation, trilinear_interpolation); }2.3 打包C扩展供Python调用创建并编写setup.py文件然后ctrl shift ~ 打开终端命令行(最好默认的改为cmd)。from setuptools import setup from torch.utils.cpp_extension import CppExtension, BuildExtension import platform # 未避免一些问题确保C桥梁文件按 C17 编译 extra_compile_args [/std:c17, /EHsc] if platform.system() Windows else [-stdc17] setup( namecppcuda_tutorial, version1.0, authorkwea123, author_emailkwea123gmail.com, descriptioncppcuda example, # cppcuda_tutorial是前面name属性值sources(可填多个值)是C中自定义的函数的别名(一般设置和自定义函数名相同) ext_modules[CppExtension(cppcuda_tutorial, sources[interpolation.cpp], extra_compile_argsextra_compile_args)], cmdclass{build_ext: BuildExtension} )在python文件中打开或者conda activate 环境名切换到c_cpp_properties.json文件设置引用的python环境中。检查并将 pip更新到最新版本。然后在项目文件夹下如下pip命令。注我使用下面命令就提示我环境中没有torch包但是我有就只是使用禁止隔离的pip命令。pip install --no-build-isolation . 或者pip install .出现Successfully installed cppcuda-tutorial-1.0说明安装成功了按我配置应该是没问题的因为我是换了个新环境重新配置的。如果从0新建的文件记得安装cuda(我的是11.8)和相应版本的torch库。2.4 创建python文件调用自定义的C函数# 一定要先引入torch库再引用自定义的c函数库 import torch import cppcuda_tutorial feats torch.randint(0,10,size (3,4)) point torch.zeros(4,3) output cppcuda_tutorial.trilinear_interpolation(feats, point) print(output)注cppcuda_tutorial库名如下图库名调用函数解析如下图实际调用的函数名也就是setup文件中的sources里面等同调用函数解析三、连通CUDA3.1 编写cu文件实现具体功能这里依旧是示例这里先返回一下参数。#include torch/extension.h torch::Tensor trilinear_fw_cu( torch::Tensor feats, torch::Tensor points ){ return feats; }3.2 配置自定义C库和修改桥梁新建文件夹include,新建.h后缀文件我的是utils.h#include torch/extension.h // 以下宏用于检查张量是否在 CUDA 上、是否连续 #define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x must be a CUDA tensor) #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x must be contiguous) #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) // 下面是cu文件中的函数用以声明方便c引用 torch::Tensor trilinear_fw_cu( torch::Tensor feats, torch::Tensor points );修改桥梁引入自定义的库在函数内部检查变量是否来自GPU否拒绝是通过然后返回cu文件中需要的函数的结果。# include torch/extension.h # include utils.h // 根据个人使用cuda的需求自定义的函数充当桥梁转接的 torch::Tensor trilinear_interpolation( torch::Tensor feats, torch::Tensor point ){ CHECK_INPUT(feats); CHECK_INPUT(point); // 调用前向传播的cuda函数 return trilinear_fw_cu(feats,point); } // 固定结构搭配 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){ // 别名setup.py文件加入环境后py文件的可以调用的函数名 自定义的函数名 m.def(trilinear_interpolation, trilinear_interpolation); }注这里那两个CHECK可能报错还有cu文件函数未识别没关系的setup文件还是可以正常编译的。如下图只能说是一个大无语。3.3 修改setup并重新编译还是在禁用隔离的情况下编译 pip install --no-build-isolation .注意里面从CppExtension改为CUDAExtension因为最后要调的是cuda文件import glob import os.path as osp import platform from setuptools import setup from torch.utils.cpp_extension import CUDAExtension, BuildExtension # 未避免一些问题确保C桥梁文件按 C17 编译 extra_compile_args [/std:c17, /EHsc] if platform.system() Windows else [-stdc17] ROOT_DIR osp.dirname(osp.abspath(__file__)) include_dirs [osp.join(ROOT_DIR, include)] # 不需要手动指定c和cuda文件 sources glob.glob(*.cpp) glob.glob(*.cu) setup( namecppcuda_tutorial, version1.0, authorkwea123, author_emailkwea123gmail.com, descriptioncppcuda_tutorial, long_descriptioncppcuda_tutorial, ext_modules[ CUDAExtension( namecppcuda_tutorial, sourcessources, include_dirsinclude_dirs, extra_compile_args{ cxx: [-O2], nvcc: [-O2] } ) ], cmdclass{ build_ext: BuildExtension } )3.4 运行py文件调用cu文件中的自定义函数和之前的一样记得将数据送到CPU上# 一定要先引入torch库再引用自定义的c函数库 import torch import cppcuda_tutorial if __name__ __main__: # feats torch.randint(0,10,size (3,4)) # feats must be a CUDA tensor # 因为要在cu文件也就是GPU中计算所以要先送到GPU上 feats torch.randint(0,10,size (3,4)).to(cuda) point torch.zeros(4,3).to(cuda) output cppcuda_tutorial.trilinear_interpolation(feats, point) print(output)结果如下图运行结果图四、添加反向传播三线性插值对输入特征feats是线性操作因此反向传播比较直接。损失对每个顶点特征的梯度等于损失对输出的梯度乘以该顶点对应的插值权重。这里设定不计算points的梯度只计算feats的梯度4.1 interpolation_kernel.cu 文件添加新增参数dL_dfeat_interp是 PyTorch 自动求导系统传入的梯度不是我们手动计算的返回值dL_feats是我们通过反向 kernel 计算并返回的最终成为feats.grad// 反向传播 template typename scalar_t __global__ void trilinear_bw_kernel( // 接收形参 const torch::PackedTensorAccessor32scalar_t, 2,torch::RestrictPtrTraits dL_dfeat_interp, const torch::PackedTensorAccessor32scalar_t, 3,torch::RestrictPtrTraits feats, const torch::PackedTensorAccessor32scalar_t, 2, torch::RestrictPtrTraits points, torch::PackedTensorAccessor32scalar_t, 3, torch::RestrictPtrTraits dL_feats ){ // 计算线程编号坐标 const int n blockIdx.x * blockDim.x threadIdx.x; const int f blockIdx.y * blockDim.y threadIdx.y; // 检查当前线程是否需要运算 if (n feats.size(0) || f feats.size(2)) return; // 定义 权重u,v,w 因为point是 N,3 // 还需要对points值做归一化使其值属于【01】 // 因为假定值在-1~1区间所以归一化为x 1/ 2 const scalar_t u (points[n][0] 1) / 2; const scalar_t v (points[n][1] 1) / 2; const scalar_t w (points[n][2] 1) / 2; const scalar_t a (1-v)*(1-w); const scalar_t b (1-v)*w; const scalar_t c v*(1-w); const scalar_t d 1-a-b-c; //转到对参数feats得梯度 dL_feats[n][0][f] (1-u)*a*dL_dfeat_interp[n][f]; dL_feats[n][1][f] (1-u)*b*dL_dfeat_interp[n][f]; dL_feats[n][2][f] (1-u)*c*dL_dfeat_interp[n][f]; dL_feats[n][3][f] (1-u)*d*dL_dfeat_interp[n][f]; dL_feats[n][4][f] u*a*dL_dfeat_interp[n][f]; dL_feats[n][5][f] u*b*dL_dfeat_interp[n][f]; dL_feats[n][6][f] u*c*dL_dfeat_interp[n][f]; dL_feats[n][7][f] u*d*dL_dfeat_interp[n][f]; } torch::Tensor trilinear_bw_cu( // 损失函数 L 对前向输出 feat_interp 的梯度 torch::Tensor dL_dfeat_interp, // 维度 N, 8, F torch::Tensor feats, // 维度 N, 3 torch::Tensor points ){ // feats 形状为 (N, 8, F)其中 8 代表三线性插值所需的八个邻近顶点F 为特征维度。、 // points 形状为 (N, 3)每个样本的三个分量分别表示在体素局部坐标系中的归一化权重范围 [0,1]分别对应 x, y, z 方向。 //输出 dL_feat 形状为 (N, 8, F)通过对八个顶点的特征损失 const int64_t N feats.size(0), F feats.size(2); // 维度 N, 8, F torch::Tensor dL_feats torch::zeros({N, 8, F}, feats.options()); const dim3 threads(16, 16); // 定义线程块Block大小 const dim3 blocks((N threads.x - 1) / threads.x, (F threads.y - 1) / threads.y); AT_DISPATCH_FLOATING_TYPES(feats.scalar_type(), trilinear_bw_cu,[] { // scalar_t 表示形状不定 trilinear_bw_kernelscalar_tblocks, threads( //packed_accessor32 将Tensor形态转换非Tensor类型参数直接传入即可。 // 参一是类型 参二是维度数 参三是一般都是torch::RestrictPtrTraits() dL_dfeat_interp.packed_accessor32scalar_t, 2, torch::RestrictPtrTraits(), feats.packed_accessor32scalar_t, 3, torch::RestrictPtrTraits(), points.packed_accessor32scalar_t, 2, torch::RestrictPtrTraits(), // 一般设定为前三个为参数最后一个是设置的接收数 dL_feats.packed_accessor32scalar_t, 3, torch::RestrictPtrTraits() ); }); return dL_feats; }4.2 interpolation.cpp 文件添加torch::Tensor trilinear_interpolation_bw( torch::Tensor dL_dfeat_interp, torch::Tensor feats, torch::Tensor points ){ CHECK_INPUT(dL_dfeat_interp); CHECK_INPUT(feats); CHECK_INPUT(points); // 调用反向传播得cuda函数 return trilinear_bw_cu(dL_dfeat_interp, feats, points); } // 固定结构搭配 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){ // 别名setup.py文件加入环境后py文件的可以调用的函数名 自定义的函数名 m.def(trilinear_interpolation_fw, trilinear_interpolation_fw); m.def(trilinear_interpolation_bw, trilinear_interpolation_bw); }4.3 utils.h 文件添加torch::Tensor trilinear_bw_cu( torch::Tensor dL_dfeat_interp, torch::Tensor feats, torch::Tensor points );4.4 测试文件添加为了让自定义 CUDA 算子参与 PyTorch 自动微分需要继承torch.autograd.Function并实现forward和backwardclass Trilinear_interpolation_cuda(torch.autograd.Function): staticmethod def forward(ctx, feats, points): feat_interp cppcuda_tutorial.trilinear_interpolation_fw(feats, points) # 保存到反向传播使用 ctx.save_for_backward(feats, points) return feat_interp staticmethod # dL_dfeat_interp对应feat_interp前向传播返回几个数反向传播就有几个形参 def backward(ctx, dL_dfeat_interp): feats, points ctx.saved_tensors dL_dfeats cppcuda_tutorial.trilinear_interpolation_bw(dL_dfeat_interp.contiguous(),feats,points) # 因为调用传入了几个参数需要计算梯度这里只有dL_dfeat_interp一个 # 需要计算梯度得就正常返回不需要的就返回None return dL_dfeats, None对比测试N 65536; F 256 rand torch.rand(N, 8, F, devicecuda) feats rand.clone().requires_grad_() feats2 rand.clone().requires_grad_() points torch.rand(N, 3, devicecuda)*2-1 t time.time() # 用到了反向传播所以同时需要前向和反向传播需要集中调用 out_cuda Trilinear_interpolation_cuda.apply(feats2, points) torch.cuda.synchronize() print( cuda time, time.time()-t, s) t time.time() out_py trilinear_interpolation_py(feats, points) torch.cuda.synchronize() print(pytorch time, time.time()-t, s) print(fw all close, torch.allclose(out_py, out_cuda)) loss out_py.sum() loss.backward() loss2 out_cuda.sum() # 执行自定义得反向传播 loss2.backward() print(bw all close, torch.allclose(feats.grad, feats2.grad))时间对比如下如所示本文内容参考YouTube博主 AI葵并结合自身环境配置和理解进行整改希望能帮到大家。