TFE：一种用于深度学习的多阶段Python嵌入式DSL

原创

已于 2022-12-07 16:20:57 修改 · 1.2k 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #tensorflow #边缘计算

于 2022-09-13 21:19:55 首次发布

TensorFlowEager是TensorFlow的一个扩展，它结合了命令式和声明式DSL的优点。默认命令式执行提供灵活性，通过function装饰器实现阶段化以提高性能。阶段化执行支持优化如并行化，并且允许在图函数中利用Python的控制流。自动微分和串行化功能方便模型训练，同时支持分布式计算。TFE的目标是提供平滑的用户体验，从本地开发到模型部署。其多阶段编程模型和自动微分特性为机器学习开发提供了新的工具。

TensorFlow Eager: A Multi-stage, Python-embedded DSL For Machine Learning

介绍和相关工作

dsl

在编程语言的观点下，DSL 可以分为命令式（Imperative）和声明式（Declarative）两种。使用命令式 DSL 编写可微分程序（基于现有自动微分软件编写的模型）就像使用命令式编程语言（比如 Python），它的性能瓶颈受制于解释器（Interpreter），并且模型的串行化比较困难。为了解决这个问题，声明式 DSL 把模型的执行与模型的定义分离，这类先定义后执行（Define-before-run）的运行库要求用户将模型划分为数据流图（Dataflow graph）的各个阶段（Stage），使得编译器优化和并行处理成为可能，从而简化部署、分布式和代码生成；但是，这也导致用户不能随意使用宿主语言的特性，而 DSL 的学习过程通常有陡峭的学习曲线；对于存在数据依赖的结构（比如控制流），这类 DSL 无法合适地描述。

理想的 DSL 应当同时具有命令式的灵活性和声明式的高效性，于是本文提出了 TensorFlow Eager。它是 TensorFlow 的一个可选扩展，只需要在程序开始时执行 tf.enable_eager_execution() 即可。TFE 默认是以命令式执行的，它另外提供了 Python 装饰器（Decorator）来追踪计算图的上下文，将建立计算图的原语操作（Primitive operations）和归属的输入输出阶段化（Staging），返回可执行的图函数（Graph function）。图函数和命令式代码共享同一个词法环境（Lexical environment），包括原语操作、计算核心和用户可见的 API 等。

在 TFE 中，用户必须手动对计算进行阶段化，可能会导致代码的重构。理想的可微分程序框架应当不需要用户干预就可以自动阶段化。一种方法（DLVM, Swift for TensorFlow, and Zygote）是将框架嵌入到预先编译好的过程式语言，把计算图提取和自动微分实现为编译器重写，但是 Python 的灵活性让 DSL 的嵌入困难重重。一些项目（AutoGraph）在 Python 的 AST 上对命令式代码进行重写，产生计算图，该技术不在本文的讨论范围内。另一种阶段化计算的选择（NVIDIA CuDNN）是融合计算核心（Fused kernel），虽然效果显著，但是泛用性差。

TFE 不是第一个提出多阶段编程模型（Multi-stage programming model）的 Python 库。JAX 是一种基于 XLA 的、为异构设备提供代码生成的 tracing-JIT 编译器，它提出了与本文类似的编程范式。MXNet 和 Gluon 也尝试在命令式计算和阶段化计算中间做插值，但抽象的视角比本文的更高。PyTorch 实现的 staging tracer 与本文的接近。可微分编程之外，Terra（基于 Lua 嵌入的 DSL）对多阶段编程的处理比本文更公式化；OptiML（基于 Scala 嵌入的 DSL）也支持阶段化，但是不支持自动微分。DSL 之外，支持 JIT 编译的还有 Numba 和 PyPy 等。

多阶段编程涉及阶段化变换（Staging transformation）和部分评估（Partial Evaluation）。

设计理念

用户编写 Python 程序的经验应当可以轻易地迁移到 TFE 的使用上，从小规模测试到在异构设备上部署模型推导的路径也需要尽可能平滑。为此，TFE 提出了三条设计原则：

特权下的命令式执行（Privilege imperative execution）。由于 Python 是命令式的，TFE 默认也是命令式执行的，阶段化执行是可选的，而且通常是不必要的。
对 Python 的无缝嵌入（Seamlessly embed into Python）。编写 TensorFlow 程序锻炼的是元编程（Metaprogramming）能力，而命令式执行让用户能够编写像宿主语言一样的程序（Pythonic），包含原生的控制流、递归、数据结构，甚至是 pdb 断点调试。
阶段化命令式代码为数据流图（Stage imperative code as dataflow graphs）。

执行模型（Execution Model）

多阶段编程

TFE 提供了两种执行算子的方式：命令式或者作为静态图的一部分。

命令式执行：默认，构造算子、立即执行。
阶段化执行：到 Python 解释器之间来回往返的开销制约了命令式执行的性能，阶段化执行不仅解除了这个限制，还能额外进行算子间并行（Inter-op parallelism）、常量折叠、缓存重用等优化。

TFE 提供了装饰器 function 来记录算子和张量流，遗憾的是，只支持 TensorFlow 中的操作，而不是任意的 Python 代码。调用装饰器返回的 Callable，就会执行装饰器产生的计算图，而不是原来的 Python 代码。计算图的运行时环境使用 C++ 编写，它会自动划分子图，分配给可用设备并尽可能尝试并行化。function 装饰器还支持基于 XLA 的代码生成，用以在 TPU 上运行。

需要注意的是，function 装饰器是 JIT tracer，在计算图上下文中（非 TensorFlow 的 Python 代码原样执行），运算返回的是计算结果的符号表示，而不是具体的值。

@tf.contrib.eager.function
def add_noise():
	eye = tf.eye(5)
	randn = np.random.randn(5, 5)
	return eye + randn

一般情况下，每次调用 add_noise() 会返回不同的结果，但是使用装饰器以后，每次调用都会产生相同的返回值，因为 add_noise() 的上下文（比如随机化种子）以常量形式记录下来了。作为推论，如果一个函数存在副作用（比如引用了一个全局自增变量），使用 function 装饰器会破坏语义等价性。

此外，由于装饰器基于 trac