在训练模型时,尤其是使用 GPU 或 TPU 等硬件加速器时,加速器可能会在等待 CPU 准备下一批数据时处于闲置状态。反之,CPU 也可能在等待加速器处理完当前批次时等待。这种“停-等”模式会引入延迟并导致硬件资源利用不足,从而减慢整体训练过程。考虑一个未经优化的简化训练循环时间线:digraph G { rankdir=LR; node [shape=box, style=filled, fontname="sans-serif"]; subgraph cluster_cpu { label = "CPU 时间"; style=filled; color="#e9ecef"; CPU_Prep_N [label="准备批次 N", fillcolor="#a5d8ff"]; CPU_Prep_N1 [label="准备批次 N+1", fillcolor="#a5d8ff"]; CPU_Prep_N2 [label="准备批次 N+2", fillcolor="#a5d8ff"]; CPU_Prep_N -> CPU_Prep_N1 -> CPU_Prep_N2 [style=invis]; // 保持对齐 } subgraph cluster_gpu { label = "GPU 时间"; style=filled; color="#e9ecef"; GPU_Train_N [label="训练批次 N", fillcolor="#b2f2bb"]; GPU_Train_N1 [label="训练批次 N+1", fillcolor="#b2f2bb"]; GPU_Train_N2 [label="训练批次 N+2", fillcolor="#b2f2bb"]; GPU_Train_N -> GPU_Train_N1 -> GPU_Train_N2 [style=invis]; // 保持对齐 } CPU_Prep_N -> GPU_Train_N [label="数据就绪"]; GPU_Train_N -> CPU_Prep_N1 [label="GPU 闲置", style=dashed, color="#adb5bd"]; CPU_Prep_N1 -> GPU_Train_N1 [label="数据就绪"]; GPU_Train_N1 -> CPU_Prep_N2 [label="GPU 闲置", style=dashed, color="#adb5bd"]; CPU_Prep_N2 -> GPU_Train_N2 [label="数据就绪"]; // 添加不可见节点和边以进行时间对齐 Start -> CPU_Prep_N [style=invis]; Start -> Idle_GPU_1 [style=invis]; CPU_Prep_N -> Idle_GPU_1 [style=invis]; Idle_GPU_1 -> GPU_Train_N [style=invis]; GPU_Train_N -> Idle_CPU_1 [style=invis]; CPU_Prep_N1 -> Idle_CPU_1 [style=invis]; Idle_CPU_1 -> CPU_Prep_N1 [style=invis]; GPU_Train_N1 -> Idle_CPU_2 [style=invis]; CPU_Prep_N2 -> Idle_CPU_2 [style=invis]; Idle_CPU_2 -> CPU_Prep_N2 [style=invis]; {rank=same; Start; Idle_GPU_1; Idle_CPU_1; Idle_CPU_2;} {rank=same; CPU_Prep_N; GPU_Train_N;} {rank=same; CPU_Prep_N1; GPU_Train_N1;} {rank=same; CPU_Prep_N2; GPU_Train_N2;} Start [shape=point]; Idle_GPU_1 [label="GPU 闲置", shape=plaintext, fontcolor="#adb5bd"]; Idle_CPU_1 [label="CPU 闲置", shape=plaintext, fontcolor="#adb5bd"]; Idle_CPU_2 [label="CPU 闲置", shape=plaintext, fontcolor="#adb5bd"]; }未预取的执行时间线。请注意 CPU 和 GPU 上的闲置期,因为它们相互等待。tf.data API 为此问题提供了一个简单解决方案:prefetch() 变换。预取使预处理和模型执行步骤重叠。当模型在 GPU 上使用批次 N 执行训练步骤(前向和后向传播)时,输入管道会使用 CPU 读取并预处理批次 N+1 的数据。这通过在管道末尾添加 dataset.prefetch(buffer_size) 来实现。buffer_size 参数指定将预取的最大元素数量(或批次,如果应用于 batch() 之后)。import tensorflow as tf # 假设 'dataset' 是初始加载后的 tf.data.Dataset 对象 # 应用如 map、shuffle、batch 等变换 dataset = dataset.map(preprocess_function, num_parallel_calls=tf.data.AUTOTUNE) dataset = dataset.shuffle(buffer_size=1000) dataset = dataset.batch(batch_size=32) # 在末尾添加预取 dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE) # 现在数据集已准备好传递给 model.fit() # model.fit(dataset, epochs=10)选择 buffer_size手动确定最佳 buffer_size 可能很复杂。它取决于预处理与模型执行所需的时间、可用内存和系统配置等因素。缓冲区过小可能无法完全隐藏数据准备延迟,而缓冲区过大可能会消耗过多内存。幸运的是,TensorFlow 提供了 tf.data.AUTOTUNE(如上例所示)。当您设置 buffer_size=tf.data.AUTOTUNE 时,tf.data 运行时会动态调整该值,试图使用保持加速器处于工作状态所需的最小缓冲区大小,同时兼顾可用内存。在大多数情况下,使用 tf.data.AUTOTUNE 是推荐的方法。prefetch() 的放置位置为了达到最佳效果,prefetch() 通常应是添加到数据集管道的最后一个变换。这确保所有前面的操作(加载、映射、混洗、批处理)都异步执行并与模型的训练步骤重叠。使用预取后,执行时间线看起来更有效率:digraph G { rankdir=LR; node [shape=box, style=filled, fontname="sans-serif"]; subgraph cluster_cpu { label = "CPU 时间"; style=filled; color="#e9ecef"; CPU_Prep_N [label="准备批次 N", fillcolor="#a5d8ff"]; CPU_Prep_N1 [label="准备批次 N+1", fillcolor="#a5d8ff"]; CPU_Prep_N2 [label="准备批次 N+2", fillcolor="#a5d8ff"]; CPU_Prep_N3 [label="准备批次 N+3", fillcolor="#a5d8ff"]; CPU_Prep_N -> CPU_Prep_N1 -> CPU_Prep_N2 -> CPU_Prep_N3 [style=invis]; // 保持对齐 } subgraph cluster_gpu { label = "GPU 时间"; style=filled; color="#e9ecef"; GPU_Idle_Start [label="初始等待", shape=plaintext, fontcolor="#adb5bd"]; GPU_Train_N [label="训练批次 N", fillcolor="#b2f2bb"]; GPU_Train_N1 [label="训练批次 N+1", fillcolor="#b2f2bb"]; GPU_Train_N2 [label="训练批次 N+2", fillcolor="#b2f2bb"]; GPU_Idle_Start -> GPU_Train_N -> GPU_Train_N1 -> GPU_Train_N2 [style=invis]; // 保持对齐 } // 显示数据流和依赖关系的箭头 CPU_Prep_N -> GPU_Train_N [label="数据就绪"]; CPU_Prep_N1 -> GPU_Train_N1 [label="数据就绪"]; CPU_Prep_N2 -> GPU_Train_N2 [label="数据就绪"]; // 使用相对定位和不可见边指示重叠 Start -> CPU_Prep_N [style=invis]; Start -> GPU_Idle_Start [style=invis]; CPU_Prep_N -> GPU_Idle_Start [style=invis]; // CPU 在 GPU 开始第一次实际工作前完成 // 对齐重叠任务的开始时间 {rank=same; CPU_Prep_N1; GPU_Train_N;} {rank=same; CPU_Prep_N2; GPU_Train_N1;} {rank=same; CPU_Prep_N3; GPU_Train_N2;} Start [shape=point]; } 使用预取的执行时间线。当 GPU 忙于训练当前批次时,CPU 准备下一个批次,这显著减少了闲置时间。通过将 .prefetch(tf.data.AUTOTUNE) 作为输入管道的最后一步,您使 TensorFlow 能够自动管理数据准备和模型计算之间的重叠,通常会以极小的代码改动带来训练吞吐量的显著提升。这是一种简单而高效的方法,用于优化数据管道的性能。