论文标题
使用XLA编译器的内存安全计算
Memory Safe Computations with XLA Compiler
论文作者
论文摘要
TensorFlow和Pytorch之类的软件包旨在支持线性代数操作，它们的速度和可用性决定了它们的成功。但是，通过优先考虑速度，他们经常忽略内存需求。结果，由于内存溢出，在软件设计方面方便的内存密集型算法的实现通常无法解决。记忆效率的解决方案需要在计算框架之外具有重要逻辑的复杂编程方法。这会损害这种算法的采用和使用。为了解决这个问题，我们开发了一个XLA编译器扩展程序，该扩展程序根据用户指定的内存限制来调整算法的计算数据流表示。我们表明，可以在单个设备上以更大的规模运行K-nearest的邻居和稀疏的高斯过程回归方法，而标准实现将失败。我们的方法可以更好地利用硬件资源。我们认为，进一步关注在编译器级别上消除内存约束将扩大将来可以开发的机器学习方法的范围。
Software packages like TensorFlow and PyTorch are designed to support linear algebra operations, and their speed and usability determine their success. However, by prioritising speed, they often neglect memory requirements. As a consequence, the implementations of memory-intensive algorithms that are convenient in terms of software design can often not be run for large problems due to memory overflows. Memory-efficient solutions require complex programming approaches with significant logic outside the computational framework. This impairs the adoption and use of such algorithms. To address this, we developed an XLA compiler extension that adjusts the computational data-flow representation of an algorithm according to a user-specified memory limit. We show that k-nearest neighbour and sparse Gaussian process regression methods can be run at a much larger scale on a single device, where standard implementations would have failed. Our approach leads to better use of hardware resources. We believe that further focus on removing memory constraints at a compiler level will widen the range of machine learning methods that can be developed in the future.
