[Paper Review] 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 Techniques26 references52 citations

TL;DR

TensorFlow Eager provides an imperative front-end to TensorFlow with a JIT tracer that can stage Python functions into dataflow graphs, enabling seamless interpolation between imperative and graph execution for ML workloads.

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.

Motivation & Objective

• Motivate the need for a DSL that combines imperative usability with graph-based optimizations for ML workflows.
• Propose a Python-embedded, multi-stage programming model that supports imperative execution by default and optional graph staging.
• Describe the implementation of an opt-in tracing-based JIT that converts Python functions into executable dataflow graphs.
• Explain how state, devices, and distribution are represented to enable seamless transitions between execution modes.
• Evaluate the performance trade-offs between imperative and staged execution on representative ML models.

Proposed method

• Introduce a Python-embedded DSL layered on TensorFlow that can execute imperatively by default and trace Python functions to build dataflow graphs when staged.
• Provide a function decorator that traces Python functions to create graph functions with named inputs/outputs and a C++ dataflow runtime for execution.
• Describe automatic differentiation via a tracing-based reverse-mode AD with nested tapes for higher-order derivatives.
• Explain state management via variables and graph-based state matching for serialization and restoration.
• Discuss device and distribution support, including automatic device placement, cross-device data transfer, and a TPU/XLA-backed staged path.
• Outline the escaping and integration mechanisms (py_func, tf.init_scope) to embed imperative code within graphs and manage state creation during tracing.

Experimental results

Research questions

• RQ1 How can an imperative-friendly front-end to TensorFlow be designed to interpolate with graph-based execution without sacrificing the benefits of dataflow graphs?
• RQ2 What are the design, implementation, and performance trade-offs of a multi-stage programming model that mixes imperative execution with staged graph construction?
• RQ3 How can state, device placement, and distribution be coherently managed across both execution modes?
• RQ4 How does tracing-based automatic differentiation interact with partially staged computations and control flow in Python?

Key findings

• Imperative TensorFlow Eager can match graph performance for ResNet-50 on a single GPU when using imperative execution.
• Staged TensorFlow Eager enables TPU-accelerated training of ResNet-50 much faster than imperative TensorFlow Eager can achieve on TPUs.
• Staging yields significant speedups for models with small operations, with minimal code changes required.
• The framework provides a coherent, single API surface that is execution-mode agnostic and leverages the TensorFlow ecosystem for tooling and deployment.
• The implementation is lightweight (Python-level staging around 2000 lines, with core differentiation in Python and C) and cross-platform.