动态张量重计算
机器学习
2021-03-19 v4 编程语言
机器学习
摘要
检查点技术通过从内存中释放中间激活并在需要时重新计算，使得在受限内存预算下训练深度学习模型成为可能。当前的检查点技术离线静态规划这些重计算，并假设计算图是静态的。我们证明了一个简单的在线算法可以通过引入动态张量重计算（DTR）达到相当的性能，DTR 是一种贪心的在线检查点算法，具有可扩展性和通用性，由驱逐策略参数化，并支持动态模型。我们证明 DTR 可以在 的内存预算下以仅 个张量操作训练一个 层线性前馈网络。在模拟实验中，DTR 紧密匹配最优静态检查点的性能。我们仅通过介入张量分配和算子调用并在张量上收集轻量级元数据，就将 DTR 原型集成到了 PyTorch 中。
引用
@article{arxiv.2006.09616,
title = {Dynamic Tensor Rematerialization},
author = {Marisa Kirisame and Steven Lyubomirsky and Altan Haan and Jennifer Brennan and Mike He and Jared Roesch and Tianqi Chen and Zachary Tatlock},
journal= {arXiv preprint arXiv:2006.09616},
year = {2021}
}
备注
31 pages, 12 figures, implementation available here: https://github.com/uwsampl/dtr-prototype, OpenReview: https://openreview.net/forum?id=Vfs_2RnOD0H