[论文解读] AutoGraph: Imperative-style Coding with Graph-based Performance

Moldovan, Dan, Decker, James M+7 more|arXiv (Cornell University)|Oct 16, 2018

DNA and Biological Computing参考文献 14被引用 10

一句话总结

AutoGraph 通过使用源代码转换和分阶段编程，将惯用的 Python 代码转换为优化的、可部署的中间表示（IR），如 TensorFlow 图或 Lantern，从而在不引入运行时开销的情况下实现原生 TensorFlow 性能，同时支持复杂的控制流。它实现了无需运行时开销的原生 TensorFlow 性能，支持复杂控制流和后端无关的编译。

ABSTRACT

There is a perceived trade-off between machine learning code that is easy to write, and machine learning code that is scalable or fast to execute. In machine learning, imperative style libraries like Autograd and PyTorch are easy to write, but suffer from high interpretive overhead and are not easily deployable in production or mobile settings. Graph-based libraries like TensorFlow and Theano benefit from whole-program optimization and can be deployed broadly, but make expressing complex models more cumbersome. We describe how the use of staged programming in Python, via source code transformation, offers a midpoint between these two library design patterns, capturing the benefits of both. A key insight is to delay all type-dependent decisions until runtime, via dynamic dispatch. We instantiate these principles in AutoGraph, a software system that improves the programming experience of the TensorFlow library, and demonstrate usability improvements with no loss in performance compared to native TensorFlow graphs. We also show that our system is backend agnostic, and demonstrate targeting an alternate IR with characteristics not found in TensorFlow graphs.

研究动机与目标

• 弥合命令式机器学习框架（易于编码）与基于图的系统（高性能、可移植）之间的差距。
• 使开发者能够使用包含控制流（例如循环、条件语句）的惯用 Python 代码，同时将其编译为优化的、静态类型的中间表示（IR）。
• 消除现有命令式和基于图的机器学习框架中固有的性能与可移植性权衡。
• 通过针对多种中间表示（IR），包括 TensorFlow 图和更具表现力的 Lantern IR，实现后端无关的编译。
• 在代码转换过程中保持语义正确性，同时通过自动类型分发和最少的注解来最小化开发者的负担。

提出的方法

• 使用 Python 中的源代码转换（SCT）分析并重写命令式代码，生成更低级的中间表示（IR）。
• 通过运行时类型信息的动态分发实现分阶段编程，将依赖类型的决策推迟到执行时才确定。
• 应用静态分析检测控制流结构（例如循环、条件语句），并将它们重写为目标 IR 中等效的基于图的操作。
• 通过为不同 IR 生成代码，支持后端无关的编译，例如 TensorFlow 图和更具表现力的 Lantern IR。
• 依赖 Python 的内省和反射（例如 inspect、imp）分析函数结构，并在运行时或编译时执行转换。
• 通过推断类型和控制流来最小化所需注解，仅在目标 IR 的类型系统有要求时才需要显式提示。

实验结果

研究问题

• RQ1 是否可以自动将具有数据依赖控制流的命令式 Python 代码转换为优化的、静态类型的 IR，且不造成性能损失？
• RQ2 源代码转换和分阶段编程在多大程度上能够保持语义正确性，同时实现在机器学习框架中的高性能执行？
• RQ3 一个系统是否能够同时支持多个 IR，包括那些比 TensorFlow 图功能更丰富的 IR（例如支持可重入函数调用）？
• RQ4 该系统如何处理类型不匹配和动态分发，以避免无声的性能下降？
• RQ5 是否可以无需大量手动注解，就使转换管道具备鲁棒性和正确性？

主要发现

• AutoGraph 能够将包含复杂控制流的惯用 Python 代码成功转换为优化的 TensorFlow 图 IR，且与手写图相比无性能损失。
• 该系统实现了原生 TensorFlow 性能，证明了通过 AutoGraph 编译的命令式编码风格不会带来运行时开销。
• AutoGraph 不仅适用于 TensorFlow，还支持更丰富的 Lantern IR，后者支持 TensorFlow 图无法表达的高级功能，如可重入函数调用。
• 转换过程在大多数情况下是全自动的，几乎不需要注解，仅在目标 IR 的类型系统有要求时才需要极少的提示。
• AutoGraph 已具备生产就绪能力，并通过 tf.function 集成到 TensorFlow 2.0 中，受益于 TensorFlow 代码库的完整测试套件。
• 尽管系统具备鲁棒性，但若用户向条件语句传递 Python 布尔值而非 TensorFlow 布尔值，系统可能会无声失败，凸显了对改进调试支持的需求。