汤不热吧在深度学习训练和推理过程中，我们经常会发现 GPU 使用率不高，或者训练速度远低于预期。这往往不是 GPU 计算慢，而是因为数据加载和预处理（Data I/O）成为了瓶颈，导致 GPU 必须等待 CPU 准备好下一批数据。这种等待被称为“数据饥饿”（Data Starvation）。
本文将详细介绍 PyTorch 中最实用的数据预取（Prefetching）技巧，通过简单配置实现 CPU 数据预处理与 GPU 计算的完美异步重叠，从而大幅提高训练效率。
问题的核心：同步等待
默认情况下，如果没有足够的异步机制，当 GPU 完成对 Batch N 的计算后，它会停下来等待 CPU 完成 Batch N+1 的加载、解码和增强。在等待期间，宝贵的 GPU 资源处于空闲状态。
解决这个问题的关键机制是 PyTorch DataLoader 中的两个核心参数：num_workers 和 pin_memory。
关键技巧一：异步并行加载 (num_workers)
num_workers 参数告诉 DataLoader 使用多少个子进程来并行地执行数据加载和预处理任务。当 GPU 正在处理 Batch N 时，这些子进程已经在后台加载和处理 Batch N+1, N+2 等数据。
- num_workers=0 (默认)：所有数据加载都在主进程中完成，与 GPU 训练完全同步，效率最低。
- num_workers > 0：启用多进程并行加载，实现 CPU 预处理与 GPU 计算的初步重叠。
关键技巧二：锁定内存加速传输 (pin_memory)
即使数据在 CPU 内存中准备好了，将其从 CPU 内存（主机端）传输到 GPU 显存（设备端）也需要时间。如果数据位于标准的 pageable CPU 内存中，驱动程序必须先将其复制到临时的 page-locked（页锁定/不可分页）内存区域，然后才能执行 DMA（直接内存访问）传输到 GPU。
设置 pin_memory=True 后，PyTorch 会自动将数据加载到 page-locked 内存中。这样，数据可以直接通过 DMA 传输到 GPU，减少了一次额外的内存拷贝，并允许传输操作本身与 GPU 计算实现更好的异步重叠。
当同时使用 pin_memory=True 和 non_blocking=True（在数据转移到 GPU 时设置）时，效果最佳。
实操代码示例：对比性能提升
我们通过一个模拟了“慢速”数据加载过程的示例来对比使用和不使用预取技巧的性能差异。
环境准备
确保你已安装 PyTorch 并拥有可用的 CUDA 设备。
import torch
from torch.utils.data import Dataset, DataLoader
import time
import os
# 模拟一个数据加载很慢的Dataset
class HeavyDataset(Dataset):
def __init__(self, size=1000, sleep_time=0.005):
# 模拟 1000 个 3x224x224 的数据块
self.data = [torch.rand(3, 224, 224) for _ in range(size)]
# 模拟 CPU 预处理耗时 (例如：图像解码和复杂增强)
self.sleep_time = sleep_time
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
# 模拟 CPU 密集型操作
time.sleep(self.sleep_time)
return self.data[idx], torch.tensor(idx % 10)
# 训练模拟函数
def train_epoch(dataloader, device, name, batch_size):
# 模拟一个简单的线性模型进行计算
model = torch.nn.Linear(224*224*3, 10).to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())
start_time = time.time()
print(f"\n--- Starting {name} (Batch Size: {batch_size}) ---")
for i, (inputs, labels) in enumerate(dataloader):
# 1. 数据转移到 GPU
# 注意：使用 non_blocking=True 确保数据传输异步进行
inputs = inputs.view(inputs.size(0), -1).to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
# 2. 模拟 GPU 计算
if i % 5 == 0:
# 模拟一个相对较长的计算时间，以确保数据预取有机会发挥作用
time.sleep(0.01)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
end_time = time.time()
duration = end_time - start_time
print(f"Time taken for {name}: {duration:.2f} seconds.")
print(f"Avg batch/sec: {len(dataloader) / duration:.2f}")
return duration
# --- 运行比较 ---
if __name__ == '__main__':
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Running on device: {DEVICE}")
# 获取可用的 CPU 核心数，用于设置 num_workers
MAX_WORKERS = os.cpu_count() or 4
BATCH_SIZE = 64
DATASET = HeavyDataset(size=1024, sleep_time=0.005) # 0.005s 模拟加载时间
# 场景 A: 无优化 (同步加载)
dataloader_slow = DataLoader(
DATASET,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=0,
pin_memory=False
)
time_slow = train_epoch(dataloader_slow, DEVICE, "Scenario A: Synchronous (num_workers=0)", BATCH_SIZE)
# 场景 B: 完整优化 (异步加载 + 锁定内存)
# 推荐 num_workers 设置为 CPU 核心数 - 1 或 2，或通过实验确定最佳值
dataloader_fast = DataLoader(
DATASET,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=min(8, MAX_WORKERS),
pin_memory=True
)
time_fast = train_epoch(dataloader_fast, DEVICE, f"Scenario B: Optimized (num_workers={min(8, MAX_WORKERS)}, pin_memory=True)", BATCH_SIZE)
print(f"\nOptimization Speedup: {time_slow / time_fast:.2f}x")
运行结果分析
在上述代码中，由于我们在 HeavyDataset 中加入了 time.sleep(0.005) 来模拟慢速 CPU 预处理，场景 A 中 GPU 必须等待每一次的 0.005s 延迟。而在场景 B 中，当 GPU 正在计算时，后台的 8 个 worker 已经并发地完成了后续数据的 0.005s 延迟，并将数据存储在锁定内存中，等待 GPU 读取。最终，场景 B 的运行时间会显著缩短，通常能达到 2x 到 5x 的速度提升（具体取决于模拟的延迟和实际硬件）。
总结最佳实践
为了确保数据加载不成为瓶颈，请始终遵循以下 PyTorch DataLoader 配置的最佳实践：
- 设置 **num_workers > 0****: 根据你的 CPU 核心数和内存情况，选择一个合适的并行加载进程数（推荐 4 到 8，或 **os.cpu_count() – 2）。
- **设置 **pin_memory=True****: 启用页锁定内存，加速数据从主机到设备端的传输。
- 使用 **non_blocking=True****: 在将 Tensor 转移到 CUDA 设备时，务必使用 **.to(device, non_blocking=True)，确保数据传输操作本身是异步的，可以与 GPU 的计算指令流重叠执行。