深度网络的内存优化
机器学习
2021-04-06 v3 计算机视觉与模式识别
摘要
深度学习正缓慢但稳步地遭遇内存瓶颈。尽管顶级 GPU 中的张量计算在过去五年中增长了 32 倍,总可用内存仅增长 2.5 倍。这阻碍了研究者探索更大架构,因为训练大型网络需要更多内存来存储中间输出。本文中,我们提出 MONeT,一个自动框架,可最小化深度网络的内存占用与计算开销。MONeT 联合优化检查点调度与各类算子的实现。MONeT 能够胜过所有先前的手工调优算子以及自动检查点。MONeT 将各类 PyTorch 模型的整体内存需求降低 3 倍,伴随 9-16% 的计算开销。在相同计算成本下,MONeT 比当前最先进的自动检查点框架少需 1.2-1.8 倍内存。我们的代码见 https://github.com/utsaslab/MONeT。
引用
@article{arxiv.2010.14501,
title = {Memory Optimization for Deep Networks},
author = {Aashaka Shah and Chao-Yuan Wu and Jayashree Mohan and Vijay Chidambaram and Philipp Krähenbühl},
journal= {arXiv preprint arXiv:2010.14501},
year = {2021}
}
备注
18 pages, ICLR'21