TensorFlow 源码拆解

附录 · 由现象回到源码的实验清单

用最小实验验证 eager、trace、tape、data 和部署的心智模型

实验 1:Eager 与 Graph 的调用次数

给函数加 Python print 和 tf.print,分别在多种 shape 下调用。Python print 主要反映 trace time,tf.print 会作为图中的执行节点;对照 f.experimental_get_tracing_count() 观察缓存行为。

实验 2:设备与同步

打开 tf.debugging.set_log_device_placement(True),在 CPU/GPU context 下执行一个矩阵乘法,最后分别在不调用 .numpy() 与调用 .numpy() 的情况下测时间。实验目标是确认提交、执行完成与结果回到 Python 是不同边界。

实验 3:梯度记录

用 persistent/non-persistent tape、显式 tape.watch 与变量自动 watch 组合测试。然后把梯度计算包进另一个 tape,观察高阶梯度为什么需要保存更多中间状态。

实验 4:输入管线

为 Dataset 的 map 函数加入 sleep,比较串行、num_parallel_calls、prefetch 和 batch 的组合。用 profiler 将生产线程、GetNext 等待和 device kernel 放在同一时间线上。

实验 5:保存与加载

导出带 input_signature 的 tf.function,检查 SavedModel 的 signatures 和 concrete functions;再用 TFLite converter 转换,并按 builder → allocate → invoke 三段排查。

这些实验的共同目标,是把“API 看起来做了什么”改成“哪一个对象在什么时间、哪个线程、哪个设备上拥有控制权”。

On this page