用于DNN内存优化训练的原地激活批归一化
计算机视觉与模式识别
2018-10-29 v3
摘要
在这项工作中,我们提出原地激活批归一化(InPlace-ABN)——一种以计算高效的方式大幅减少现代深度神经网络训练内存占用的新方法。我们的方案用单个插件层替代常规使用的BatchNorm+激活层序列,从而避免侵入式框架改动,同时为现有深度学习框架提供直接适用性。我们通过丢弃中间结果并在反向传播中通过反转存储的前向结果来恢复所需信息,获得了高达50%的内存节省,仅带来少量(0.8-2%)的计算时间增加。此外,我们展示了如何将常用检查点方法变得在计算上与原地激活批归一化一样高效。在图像分类实验中,我们在ImageNet-1k上展示了与最先进方法相当的结果。在内存需求高的语义分割任务中,我们报告了COCO-Stuff、Cityscapes和Mapillary Vistas的结果,并在后者上以单尺度、单模型场景在不额外训练数据的情况下取得了新的SOTA结果。代码可在https://github.com/mapillary/inplace_abn找到。
引用
@article{arxiv.1712.02616,
title = {In-Place Activated BatchNorm for Memory-Optimized Training of DNNs},
author = {Samuel Rota Bulò and Lorenzo Porzi and Peter Kontschieder},
journal= {arXiv preprint arXiv:1712.02616},
year = {2018}
}