[论文解读] TensorFlow Eager: A Multi-Stage, Python-Embedded DSL for Machine Learning

Akshay Agrawal, Akshay Naresh Modi+9 more|arXiv (Cornell University)|Feb 27, 2019

Parallel Computing and Optimization Techniques参考文献 26被引用 52

一句话总结

TensorFlow Eager 提供一个对 TensorFlow 的即时前端，带有能够将 Python 函数阶段化为数据流图的 JIT 追踪器，从而实现对 ML 工作负载在即时模式和图执行之间的无缝插值。

ABSTRACT

TensorFlow Eager is a multi-stage, Python-embedded domain-specific language for hardware-accelerated machine learning, suitable for both interactive research and production. TensorFlow, which TensorFlow Eager extends, requires users to represent computations as dataflow graphs; this permits compiler optimizations and simplifies deployment but hinders rapid prototyping and run-time dynamism. TensorFlow Eager eliminates these usability costs without sacrificing the benefits furnished by graphs: It provides an imperative front-end to TensorFlow that executes operations immediately and a JIT tracer that translates Python functions composed of TensorFlow operations into executable dataflow graphs. TensorFlow Eager thus offers a multi-stage programming model that makes it easy to interpolate between imperative and staged execution in a single package.

研究动机与目标

• 激发对将即时易用性与基于图的优化相结合以用于 ML 工作流的 DSL 的需求。
• 提出一个嵌入 Python 的多阶段编程模型，默认支持即时执行，并可选的图阶段化。
• 描述一个可选参与的基于追踪的 JIT 实现，将 Python 函数转换为可执行的数据流图。
• 解释如何表示状态、设备和分布，以实现执行模式之间的无缝转换。
• 评估在代表性 ML 模型上即时执行和阶段化执行之间的性能权衡。

提出的方法

• 引入一个基于 TensorFlow 的 Python 嵌入式 DSL，默认可以以即时方式执行，并在阶段化时追踪 Python 函数以构建数据流图。
• 提供一个函数装饰器，用于追踪 Python 函数，创建具有命名输入/输出的图函数，以及用于执行的 C++ 数据流运行时。
• 描述通过基于追踪的反向模式自动微分（AD）实现，采用带有嵌套 Tape 的高阶导数。
• 通过变量和基于图的状态匹配实现状态管理，以进行序列化和还原。
• 讨论设备与分布的支持，包括自动设备放置、跨设备数据传输，以及基于 TPU/XLA 的阶段化路径。
• 概述逃逸与集成机制（py_func、tf.init_scope），将即时代码嵌入到图中并在追踪时管理状态创建。

实验结果

研究问题

• RQ1 如何设计一个对即时友好的 TensorFlow 前端，使其与基于图的执行插值，同时不牺牲数据流图的优势？
• RQ2 混合即时执行与阶段化图构建的多阶段编程模型在设计、实现和性能上的权衡是什么？
• RQ3 如何在两种执行模式下协同管理状态、设备放置和分布？
• RQ4 基于追踪的自动微分如何与部分阶段化的计算和 Python 内的控制流交互？

主要发现

• 在使用即时执行时，TensorFlow Eager 在单个 GPU 上的 ResNet-50 可以达到与图执行相匹配的性能。
• 阶段化的 TensorFlow Eager 使 ResNet-50 的 TPU 加速训练远快于即时 TensorFlow Eager 在 TPU 上可以达到的性能。
• 阶段化在具有较小操作的模型上带来显著的加速，且所需的代码修改很少。
• 该框架提供一个一致、单一的 API 表面，与执行模式无关，并利用 TensorFlow 生态系统进行工具链和部署。
• 实现轻量级（Python 级阶段约 2000 行，核心的差异在 Python 与 C），并且跨平台。