TensorFlow 源码拆解

3 · Eager 与 Graph:两套语义如何共存

从 tf.function tracing 到 ConcreteFunction、Graph Executor 和异步 eager

Eager 是默认交互层

普通 tf.add 调用会在 Python 侧完成输入转换,然后进入 execute.quick_execute。该函数把 op name、输出数量、输入 handle、序列化 attrs 和当前 context 交给低层 API。低层返回输出 handle,Python 再把它们包装成 Tensor。

Python op wrapper
  → convert_to_tensor / dtype / attrs
  → quick_execute
  → EagerContext::Execute / device client
  → kernel registry lookup
  → OpKernel::Compute
  → output Tensor handles

Eager 的优势是 Python 控制流、调试和错误位置自然;代价是每次调用都需要派发和可能的 Python 开销,且跨设备与异步完成时间更难直接观察。

tf.function 的本质是多态函数缓存

polymorphic_function.Function 并不是一个只存“编译结果”的装饰器。它持有 Python 函数、FunctionSpec、追踪计数、捕获变量、ConcreteFunction 集合和序列化所需的 Trackable 信息。调用时大致经历:

  1. 将位置参数、关键字参数和嵌套结构规范化;
  2. 从值、dtype、shape、Python 常量和 TypeSpec 生成 tracing type;
  3. 在缓存中查找兼容的 ConcreteFunction;
  4. 若没有命中,创建 FuncGraph 并执行一次 tracing;
  5. 将输出结构和 captures 固化到 ConcreteFunction;
  6. 后续调用按 flat signature 传入 Tensor 和 captured inputs。

因此同一个 Python 函数可能有多个 ConcreteFunction。输入 shape、Python bool、列表结构和 input_signature 都会影响缓存键;频繁 retracing 往往是签名不稳定的信号。

FuncGraph 负责“把 Python 语义翻译成图”

FuncGraph 继承 ops.Graph,增加了 outer graph、captures、变量、deferred captures 和保存错误等信息。外层 eager Tensor 若在 tracing 中被引用,不是直接把数值拷贝进图,而是通过 capture() 创建内部 placeholder,并保存外部值与内部输入的映射。

outer eager value ──capture──> FuncGraph placeholder
                                  │
                                  └─ ConcreteFunction input signature

这也是闭包变量、资源变量和 SavedModel 可导出的关键:图中只出现可序列化的输入/资源引用,运行时再把 capture 绑定回真实对象。

Graph 执行如何调度

tensorflow/core/graph/graph.h 中的 Node 与 Edge 组成有向图。executor.cc 初始化节点元信息,然后 ExecutorState 维护 ready queue;当节点输入满足依赖时,Process 准备输入并调用同步或异步 kernel,完成后把下游节点加入队列。控制边也参与就绪判断,所以副作用 op 不只是按数据依赖执行。

失败模式:Python 控制流与 Tensor 控制流

@tf.function
def f(x):
  if x > 0:       # 这里的 x 是 Tensor,不能按普通 Python bool 处理
    return x * 2
  return x - 2

AutoGraph 可能把 if 转换为 TensorFlow control-flow op;但 Python side effect、迭代器、可变容器仍会在 tracing time 发生。读代码时要先问:“这段语句发生在 trace time 还是 graph run time?”

源码锚点

  • tensorflow/python/eager/execute.py:eager op 统一派发。
  • tensorflow/python/eager/polymorphic_function/polymorphic_function.py:Function.__call__、缓存和 tracing。
  • tensorflow/python/eager/polymorphic_function/concrete_function.py:ConcreteFunction.__call__、captures、签名。
  • tensorflow/python/framework/func_graph.py:FuncGraph 与 func_graph_from_py_func。
  • tensorflow/core/graph/graph.h:Node、Edge、输入输出和控制流分类。
  • tensorflow/core/common_runtime/executor.cc:ready queue、同步/异步 kernel 调度。

On this page