1. 使用 DataLoader 的多进程加载
DataLoader(dataset, batch_size=64, num_workers=4, pin_memory=True)
2. 混合精度训练
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
with autocast():
output = model(input)
loss = criterion(output, target)
scaler.scale(loss).backward()
3. 梯度累积
当显存不够时,可以通过梯度累积模拟更大的 batch size。