5 · tf.data：把输入变成可调度的计算

Dataset 是惰性计算图

tf.data.Dataset 表面上像 Python iterable，内部更接近一棵数据变换树：map、batch、shuffle、interleave 等操作创建新的 dataset variant Tensor，并把 upstream dataset、函数和属性编码进节点。构造 pipeline 通常不读取数据；读取发生在 iterator 的 GetNext。

source Dataset
  → MapDataset(function)
  → BatchDataset(batch_size)
  → PrefetchDataset(buffer_size)
  → Iterator::GetNext
  → Tensor(s) / structure

tensorflow/python/data/ops/dataset_ops.py 维护 Dataset API、element_spec、嵌套结构与 variant tensor；C++ tensorflow/core/data 管理 iterator、dataset kernel、线程池、buffer、checkpoint 和 service。一个 map 函数可能是 eager callable，也可能被包装成 FunctionDef，由 captured_function 在迭代器线程中执行。

element_spec 描述的不只是 dtype/shape，还描述嵌套结构和可选、稀疏、组合 Tensor 的类型系统。Keras、tf.function 和分布式输入都依赖它来建立签名。

Prefetch 不是魔法加速

prefetch 让生产者在消费者训练 step 之前填充 buffer。其效果取决于：

上游 map 是否受 Python GIL 或昂贵 decode 限制；
buffer size 是否足以隐藏计算时间；
iterator 线程池与 inter-op/intra-op 线程的竞争；
device prefetch 是否把 batch 搬到目标设备；
checkpoint/重复迭代时 buffer 如何恢复。

源码中需要同时读 Dataset op 的 MakeIterator、iterator 的 GetNext 和 cancellation/cleanup，否则只看 prefetch() API 会漏掉 shutdown 路径。

Autotune 的位置

AUTOTUNE 会把并行度、buffer size 等变成 runtime 可调整参数。调优器依赖采样到的生产/消费时间和内存预算，实际是一个反馈控制问题，而不是固定常量。性能分析时应区分 input pipeline 饥饿、模型 kernel 慢和设备 copy 慢。

分布式 data service

tensorflow/core/data/service 把数据处理拆成 dispatcher、worker、任务和 consumer。客户端创建数据集处理任务，worker 执行 dataset graph 和迭代；consumer 通过 task runner 拉取数据。故障恢复、sharding、跨 trainer 共享和 checkpoint 都会影响语义。

trainer / consumer
       │ register + get task
       ▼
dispatcher ───────────► worker 1 ─► dataset graph
       └──────────────► worker 2 ─► dataset graph

源码锚点

tensorflow/python/data/ops/dataset_ops.py：Dataset API、variant、element_spec。
tensorflow/python/data/ops/iterator_ops.py：Iterator、GetNext、checkpoint 边界。
tensorflow/core/data/root_dataset.h / dataset.h：C++ Dataset/Iterator 抽象。
tensorflow/core/data/captured_function.h：在 data runtime 中执行捕获函数。
tensorflow/core/data/service/：dispatcher、worker、task、client。
tensorflow/core/data/optimization/：静态与运行时 pipeline 优化。

阅读实验

ds = tf.data.Dataset.range(8)
ds = ds.map(lambda x: x * 2, num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.batch(2).prefetch(tf.data.AUTOTUNE)
print(ds.element_spec)
for batch in ds: print(batch.numpy())