最近遇到了模型性能方面的问题,调研中关于JIT(just in time)即时编译一些知识点进行介绍:
概述
XLA(加速线性代数)是用于优化TensorFlow计算的线性代数的域特定编译器。代码位置在tensorflow/compiler.
在XLA技术之前,TensorFlow中计算图的执行是由runtime(运行时)代码驱动的:runtime负责加载计算图定义、创建计算图、计算图分区、计算图优化、分配设备、管理节点间的依赖并调度节点kernel的执行;计算图是数据部分,runtime是代码部分。在XLA出现之后,我们有了另一个选择,计算图现在可以直接被编译成目标平台的可执行代码,可以直接执行,不需要runtime代码的参与了。
XLA 利用 JIT 编译技术分析用户在运行时创建的 TensorFlow 图表,根据实际运行时维度和类型将其专门化,将多个运算融合在一起并为它们生成高效的本机代码——适用于 CPU、GPU 之类的设备和自定义加速器(例如,Google 的 TPU)。
目前XLA是实验性的。大多数使用情况在性能(加快速度或减少内存使用)方面都没有改进。
代码示例
代码来自tenorflow源码下的tensorflow\examples\tutorials\mnist\mnist_softmax_xla.py
这份代码原理和前面几篇博客类似,相通的知识点就不特别说了。
开启JIT编译
在会话级别打开JIT方法如下:
方式一,通过Session设置:
<