XLA(Accelerated Linear Algebra)-加速线性代数,是Google推出的高性能机器学习领域编译器,它可以在不更改源代码的条件下加速Tensorflow模型。
TensorFlow在设计时主要考虑灵活性和可拓展性,而灵活性往往和高性能不可兼得,因此TensorFlow的计算性能有所欠缺。XLA就是为了解决这个问题而提出的。
XLA为什么有效
-
融合可组合算子从而提高性能
XLA通过对TensorFlow运行时的计算图进行分析,将多个低级算子进行融合从而生成高效的机器码。例如:
因为上面计算图中的许多算子都是逐元素(element-wise)的计算,所以可以融合为一个element-wise的循环计算kernel中。matmul的结果加biases时为element-wise;然后对Add的结果每个进行Relu;再对Relu的结果每个元素取幂运算。
通过将这些算子融合,可以减少申请这些算子间的中间结果所占用的内存。同时因为将多个kernel融合为了一个kernel,因此减少了加载kernel的时间消耗。
因此,XLA对模型的性能提升主要来自于将多个连续的element-wise算子融合为一个算子。 -
提高内存利用率
通过对内存使用的分析与规划,原则上可以消除许多中间数据的内存占用。 -
减少模型可执行文件大小
对于移动设备场景,XLA可以减少TensorFlow模型的执行文件大小。通过AOT(ahead-of-time