TFLMS：通过图重写实现 TensorFlow 中的大模型支持
机器学习
2019-10-03 v2 人工智能
机器学习
摘要
尽管 GPU 等加速器内存有限，深度神经网络却变得越来越大，在训练时将无法适配加速器的内存限制。我们提出一种通过重写神经网络计算图来解决该问题的方法，其中插入换出（swap-out）与换入（swap-in）操作，将中间结果临时存储于 CPU 内存。具体而言，我们首先通过为图中的变量定义具体语义来修正计算图的概念。随后我们形式化地展示如何从已有图推导换出与换入操作，并给出优化该图的规则。为实现我们的方法，我们在 TensorFlow 中开发了一个名为 TFLMS 的模块。TFLMS 作为拉取请求发布于 TensorFlow 代码仓库，以贡献 TensorFlow 社区。借助 TFLMS，我们得以将 ResNet-50 和 3DUnet 的批大小分别增大 4.7 倍和 2 倍进行训练。特别是，我们能够使用尺寸为 的图像训练 3DUNet 用于图像分割，而若无 TFLMS，此前只能通过将该图像分割为更小图像来进行，这会影响精度。
引用
@article{arxiv.1807.02037,
title = {TFLMS: Large Model Support in TensorFlow by Graph Rewriting},
author = {Tung D. Le and Haruki Imai and Yasushi Negishi and Kiyokuni Kawachiya},
journal= {arXiv preprint arXiv:1807.02037},
year = {2019}
}
备注
A new version of TFLMS was published at ISMM 2019 (https://dl.acm.org/citation.cfm?id=3329984)