OccNet显存优化:用混合精度与梯度检查点训练3D占用模型的4个方案 OccNet显存优化用混合精度与梯度检查点训练3D占用模型的4个方案【免费下载链接】OccNet[ICCV 2023] OccNet: Scene as Occupancy项目地址: https://gitcode.com/gh_mirrors/oc/OccNetOccNet 是 ICCV 2023 提出的 3D 占用Scene as Occupancy模型它把自动驾驶场景建模为带语义的 3D 体素网格相比传统 3D 检测能更精细地刻画障碍物和可通行空间。但代价也很直观以 200×200 的 BEV 特征、16 层柱状查询和 17 类语义光流输出计算OccNet 训练时的显存开销极大很多同学在 24GB 的 3090/4090 上跑基线都频频 OOM。本文分享 4 个经过实践验证的OccNet 显存优化方案重点讲解混合精度训练与梯度检查点的配置方法帮你把 3D 占用模型稳稳塞进单卡显存。上图是 OccNet 的典型输出上半部分是环视相机输入下半部分是对场景的 3D 语义占用预测每个体素都被赋予道路、植被、车辆等语义类别。要训练出这样的效果先看看显存都花在哪了——ResNet50 骨干 FPN 颈部 4 层 Transformer 编码器 3D 占用头光中间激活值就能吃掉十几 GB。方案一开启混合精度训练fp16显存直接减半这是OccNet 显存优化里性价比最高的方案几乎零成本就能把显存占用砍掉约 40%50%。项目训练流程位于 mmdet_train.py其中已经内置了对 fp16 的支持读取配置中的fp16字段自动实例化Fp16OptimizerHook实现前向半精度、梯度回传与参数更新的混合精度训练。你只需在 bevformer_base_occ.py 里追加fp16 dict(loss_scaledict(init_scale512))即可启用。默认init_scale512能较好平衡数值稳定与性能若训练中频繁出现overflow警告可适当调低至 256。配合动态损失缩放OccNet 的精度几乎不受影响而显存占用立竿见影地下降。方案二用梯度检查点with_cp换显存如果你开启了 fp16 依然吃紧第二步就是梯度检查点。它的原理很简单反向传播时不再保留每一层的中间激活值而是前向重算一遍用少量计算时间换取大量显存。在 bevformer_base_occ.py 的img_backbone配置中把默认的with_cpFalse改为img_backbonedict( typeResNet, depth50, ... with_cpTrue),ResNet50 在 6 视角 1600×900 输入下的激活值非常可观开启后骨干部分显存可再省 35GB代价是训练速度下降约 10%15%属于典型的以时间换空间。如果还想更进一步可把 TransformerOcc 头中的 attention 层也换成可重算实现不过这需要改源码普通用户做到骨干这步已经够用。方案三调低批大小并控制数据加载开销在 bevformer_base_occ.py 的data配置中基线默认samples_per_gpu1、workers_per_gpu4。3D 占用模型的 batch 大小对显存呈线性影响因此保持samples_per_gpu1不要盲目加 batch先用方案一、二腾出空间若仍 OOM可把workers_per_gpu从 4 降到 2减少 CPU 预取与 pinned memory 占用检查data.train.pipeline中的图像尺寸必要时在PadMultiViewImage前缩小img_scale从 (1600, 900) 降到 (1280, 720) 可进一步压缩输入显存。这一招与混合精度、梯度检查点叠加使用OccNet 在 24GB 显存上跑完整训练不再是奢望。方案四多卡分布式训练摊薄单卡压力最后别忘了项目自带的分布式训练脚本 dist_train.sh它基于tools/train.py实现多卡并行./tools/dist_train.sh projects/configs/bevformer/bevformer_base_occ.py 8把 8 换成你实际拥有的 GPU 数量即可。在多卡场景下每张卡的 batch 独立计算单卡显存压力被天然摊薄再叠加前面的 fp16 与 with_cp甚至可以在 2 张 16GB 的卡上完成 24 epoch 的完整训练。需要注意的是分布式模式下mmdet_train.py会使用MMDistributedDataParallel如果开启梯度检查点导致某些参数未参与梯度计算记得在配置中设置find_unused_parametersTrue。总结4 个方案怎么组合最优方案显存收益速度影响难度混合精度 fp16减少约 40%50%更快⭐梯度检查点 with_cp骨干省 35GB慢 10%15%⭐⭐调低 batch 与分辨率线性下降略慢⭐多卡分布式训练按卡数摊薄更快⭐⭐⭐推荐顺序先开 fp16再开 with_cp仍不够就降分辨率最后上多卡。这套组合拳下来OccNet 这类 3D 占用模型在单张消费级显卡上也能顺利训练。动手之前记得先git clone https://gitcode.com/gh_mirrors/oc/OccNet获取完整代码与 getting_started.md 中的环境搭建指引再用tools/train.py启动你的实验吧【免费下载链接】OccNet[ICCV 2023] OccNet: Scene as Occupancy项目地址: https://gitcode.com/gh_mirrors/oc/OccNet创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考