Bev pool 加速(1): torch.autograd.Function的使用
文章目录
- 1. 介绍
- 2 torch.autograd.Function
- 2.1 自动微分
- 2.2 autograd.Function的使用
- 2.3 案例
- 2.3.1 forward 方法
- 2.3.2 setup_context方法
- 2.3.3 backward方法
- 2.3.2.1 backward 讲解
- 2.3.2.2 梯度计算公式
- 2.3.4 验证梯度计算
- 2.3.5 使用自定义的backward
- 3. 完整代码
1. 介绍
在做Bev项目时，我们会发对于一些
计算量非常大的算子，通常会通过CUDA进行加速实现，将cuda计算封装到一个动态库.so，然后提供给python调用，相对于原先的python加速能显著的提升计算效率。比如bev_pool就是通过这种方式实现离线计算和cuda并行，从而速度提升了几个量级
bev_pool是Bev算法的关键的运算瓶颈，针对该问题工程上做了大量的性能优化，包括离线运算和cuda并行计算，从而实现bev_pool的计算加速。代码示例如下
bevfusion/mmdet3d/ops/bev_pool/bev_pool.py
class QuickCumsumCuda(torch.autograd.Function