TFLMS: Large Model Support in TensorFlow by Graph Rewriting
This addresses memory constraints for deep learning practitioners using TensorFlow, enabling training of larger models or with larger batch sizes, though it is an incremental improvement on existing swapping techniques.
The paper tackles the problem of training large neural networks on GPUs with limited memory by proposing TFLMS, a TensorFlow module that rewrites computational graphs to swap intermediate results to CPU memory, enabling training of ResNet-50 and 3DUnet with 4.7x and 2x larger batch sizes, respectively, and allowing 3DUNet to use full-size images for segmentation without accuracy loss.
While accelerators such as GPUs have limited memory, deep neural networks are becoming larger and will not fit with the memory limitation of accelerators for training. We propose an approach to tackle this problem by rewriting the computational graph of a neural network, in which swap-out and swap-in operations are inserted to temporarily store intermediate results on CPU memory. In particular, we first revise the concept of a computational graph by defining a concrete semantics for variables in a graph. We then formally show how to derive swap-out and swap-in operations from an existing graph and present rules to optimize the graph. To realize our approach, we developed a module in TensorFlow, named TFLMS. TFLMS is published as a pull request in the TensorFlow repository for contributing to the TensorFlow community. With TFLMS, we were able to train ResNet-50 and 3DUnet with 4.7x and 2x larger batch size, respectively. In particular, we were able to train 3DUNet using images of size of $192^3$ for image segmentation, which, without TFLMS, had been done only by dividing the images to smaller images, which affects the accuracy.