DenseNet 的内存高效实现
计算机视觉与模式识别
2017-07-24 v1
摘要
DenseNet 架构因特征重用而具有极高的计算效率。然而,简单的 DenseNet 实现可能需要大量 GPU 内存:如果管理不当,预激活批归一化和连续卷积操作产生的特征图会随网络深度呈二次增长。在本技术报告中,我们介绍了减少 DenseNet 在训练期间内存消耗的策略。通过策略性地使用共享内存分配,我们将存储特征图的内存开销从二次降至线性。消除了 GPU 内存瓶颈后,现在可以训练极深的 DenseNet。单张 GPU 上可训练的模型参数量从 4M 提升至 14M。此前无法训练的 264 层 DenseNet(73M 参数),现在可以在配备 8 块 NVIDIA Tesla M40 GPU 的单工作站上训练。在 ImageNet ILSVRC 分类数据集上,这个大型 DenseNet 获得了 20.26% 的 SOTA 单次裁剪 top-1 错误率。
引用
@article{arxiv.1707.06990,
title = {Memory-Efficient Implementation of DenseNets},
author = {Geoff Pleiss and Danlu Chen and Gao Huang and Tongcheng Li and Laurens van der Maaten and Kilian Q. Weinberger},
journal= {arXiv preprint arXiv:1707.06990},
year = {2017}
}
备注
Technical report