2 个月前

原位激活批量归一化用于优化DNN训练的内存使用

Samuel Rota Bulò; Lorenzo Porzi; Peter Kontschieder
原位激活批量归一化用于优化DNN训练的内存使用
摘要

在本研究中,我们提出了一种新的方法——原位激活批量归一化(In-Place Activated Batch Normalization, InPlace-ABN),该方法能够在计算效率较高的情况下显著减少现代深度神经网络的训练内存占用。我们的解决方案以单一插件层替代了传统使用的批量归一化(BatchNorm)+ 激活函数(Activation)层的组合,从而避免了对框架进行侵入式修改,同时为现有的深度学习框架提供了直接的应用可能性。通过丢弃中间结果并在反向传播过程中通过存储的前向结果的逆运算恢复所需信息,我们最多可节省50%的内存,而计算时间仅增加0.8%-2%。此外,我们还展示了如何使常用的检查点方法在计算效率上与InPlace-ABN相当。在图像分类实验中,我们在ImageNet-1k数据集上展示了与现有最先进方法相当的结果。对于内存需求较高的语义分割任务,我们在COCO-Stuff、Cityscapes和Mapillary Vistas数据集上报告了实验结果,并在单尺度、单模型的情况下,在后者数据集上取得了新的最先进结果,而无需额外的训练数据。代码可在https://github.com/mapillary/inplace_abn 获取。