算子加速(3):自定义cuda扩展
需要自定义某个层,或有时候用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++文件
,该文件定义了将从Pyt