6 · Keras 训练系统：高层 API 如何落到底层 step

Keras 是 TensorFlow 的编排层

Layer 负责可复用计算与权重追踪，Model 在其上增加训练、评估、预测、保存与指标；Functional / Sequential 是不同的模型构造方式，但最终都汇入 Layer call 与训练 loop。Keras 不是另一个 runtime，它把 TensorFlow Tensor、Variable、Function、GradientTape 和 Dataset 编排成用户可用的生命周期。

Layer.__call__
  → build(input_shape) once
  → call(inputs, training, mask)
  → Trackable weights / losses / updates

Model.fit
  → data_adapter
  → train_function (possibly tf.function)
  → train_step
  → GradientTape + optimizer.apply_gradients

Layer 的关键机制：追踪而不只是封装

tensorflow/python/keras/engine/base_layer.py 的 Layer 继承 Trackable 体系。赋值给 layer 的子层、Variable、metric 和 regularizer 会被发现并纳入权重、checkpoint 与序列化。build() 延迟创建权重，使 layer 可以从首次输入推断 shape；但这也要求 build 过程在 tracing 和分布式环境中保持可控。

Layer.__call__ 还会处理输入结构、mask、training 参数、autocast、活动正则化和 symbolically constructed tensors。源码阅读时应把“用户写的 call”和“外围包装器的 side effects”分开。

compile 只是配置，不是立即训练

Model.compile 保存 optimizer、loss、metrics、run_eagerly、steps_per_execution 等配置，并清空可能已经构建的 train/test/predict function。真正的 train_function 通常在第一次 fit 时通过 make_train_function 懒构建；它可以被 tf.function 包装，也可以用 run_eagerly=True 保留 Python 调试语义。

train_step 是最短的自定义入口

class MyModel(tf.keras.Model):
  def train_step(self, data):
    x, y = data
    with tf.GradientTape() as tape:
      y_pred = self(x, training=True)
      loss = self.compute_loss(x, y, y_pred)
    grads = tape.gradient(loss, self.trainable_variables)
    self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
    return {m.name: m.result() for m in self.metrics}

这段代码把第 4 章的 Tape、第 5 章的 Dataset 和当前章节的 Model 状态连接起来。fit 的复杂性主要来自 data adapter、epoch/step/callback、分布式包装和函数缓存，而不是来自这 10 行数学逻辑。

数据适配器与多输入

engine/data_adapter.py 把 NumPy、Tensor、Dataset、generator 等输入统一成可迭代 batch，并推导 cardinality、batch size、partial batch、sample weight 与 validation 数据。模型的输入结构必须同时满足 Python nesting、TypeSpec 与 layer 的输入签名；很多“训练时报结构错误”其实发生在 adapter 到 train_step 的边界。

保存与 Trackable 图

Keras 保存不只保存权重。模型会把 layer、variable、optimizer slot、函数和配置组织成 Trackable object graph；SavedModel 则将可调用签名和 ConcreteFunction 一起导出。恢复时，名字和对象路径、slot variables、自定义对象注册都会参与匹配。

源码锚点

• tensorflow/python/keras/engine/base_layer.py：Layer call、build、权重追踪。
• tensorflow/python/keras/engine/training.py：Model、compile、fit、train_step、保存。
• tensorflow/python/keras/engine/data_adapter.py：输入统一与 batch 语义。
• tensorflow/python/keras/engine/functional.py：符号图模型。
• tensorflow/python/trackable/：对象图、依赖追踪、checkpoint。
• tensorflow/python/keras/saving/：SavedModel/配置/权重序列化。