文章目录
- 1.CUDA 扩展介绍
- 1.1 CUDA 检查
- 1.2 项目结构
- 2 简单案例
- 2.1 c++ 代码
- 2.2 cu 核函数
- 2.3 setup.py文件
- 2.4 python文件
- 3 进阶案例
- 3.1 c++ 代码实现
- 3.2 cuda 核函数的实现
- 3.2.1 forward 方法
- 3.2.2 backward 方法
- 3.3 setup.py
- 3.4 python 调用
- 4.完整代码
需要自定义某个层，或有时候用c++实现你的操作(c++扩展)可能会更好:
- 例如：需要实现一个新型的激活函数
- 例如: bevfusion用cuda实现bevpool加速
自定义扩展的步骤
- (1) 首先用纯pytorch和python 实现我们所需的功能，看看效果再决定要不要进一步优化
- (2) 明确优化方向，用C++ (或CUDA) 重写部分代码
- (3) 用纯C++编写它
- (4) 将模型的部分移动到CUDA内核来进一步加速，以便从GPU提供的大规模并行计算中获益
在前面算子加速(2):自定义c++扩展，第(4)步cuda核函数来实现加速，是没有用到的。如果想进一步提升速度的话，那么就需要通过CUDA 核函数来实现部分功能，也就是实现 cuda扩展。
1.CUDA 扩展介绍
编写CUDA扩展的一般策略是首先编写一个
C++文件
国际智能手机供应链展览会)
![[数据集][目标检测]高铁受电弓检测数据集VOC+YOLO格式1245张2类别](http://pic.xiahunao.cn/[数据集][目标检测]高铁受电弓检测数据集VOC+YOLO格式1245张2类别)