PyTorch多进程数据加载时，在__getitem__()中直接调用.cuda()会引发GPU内存泄漏——子进程无法正确释放CUDA上下文，导致empty_cache()失效；正确做法是CPU端完成数据加载与预处理，仅在主训练循环中统一、非阻塞地迁移至GPU。
pytorch多进程数据加载时，在`__getitem__()`中直接调用`.cuda()`会引发gpu内存泄漏——子进程无法正确释放cuda上下文，导致`empty_cache()`失效；正确做法是cpu端完成数据加载与预处理，仅在主训练循环中统一、非阻塞地迁移至gpu。
在PyTorch分布式数据加载场景中，GPU内存“只增不减”是一个典型且易被忽视的问题。其根本原因并非显存碎片或梯度累积，而是CUDA张量生命周期与多进程模型的冲突：当 DataLoader 的 num_workers > 0 时，每个 worker 子进程会独立初始化 CUDA 上下文（即使未显式调用 torch.cuda.*），而子进程中创建的 CUDA 张量无法被主进程的 torch.cuda.empty_cache() 清理，且子进程退出时 CUDA 内存往往不会立即归还驱动，从而表现为 VRAM 持续攀升。
✅ 正确实践：CPU 加载 + 主进程 GPU 迁移
首先，重构 CustomDataset.__getitem__()，确保所有操作严格运行在 CPU 上：
class CustomDataset(Dataset):
def __init__(self, data_paths):
self.data_paths = data_paths
def __len__(self):
return len(self.data_paths)
def __getitem__(self, idx):
# ✅ 全部在 CPU 完成：加载、类型转换、转为 Tensor
image = np.load(self.data_paths[idx]['image']).astype(np.float32)
label = np.load(self.data_paths[idx]['label']).astype(np.int64)
# 使用 from_numpy 避免额外拷贝，且不触发 CUDA
return torch.from_numpy(image), torch.from_numpy(label)
然后，在训练主循环中，在 DataLoader 返回 batch 后，使用 non_blocking=True 高效迁移至 GPU：
for epoch in range(10):
for batch in dataloader:
images, labels = batch
# ✅ 主进程内迁移：支持异步传输，提升吞吐
images = images.cuda(non_blocking=True)
labels = labels.cuda(non_blocking=True)
output = images.mean()
loss = output.sum()
loss.backward()
# ✅ 清理计算图引用（非必需但推荐）
del images, labels, output, loss
# ⚠️ 注意：此处 empty_cache() 通常无需每步调用，见下方说明
? 关键配置增强稳定性
-
启用
persistent_workers=True（PyTorch ≥ 1.7）：复用 worker 进程，避免反复 fork/cuda 初始化开销，显著缓解内存缓慢增长：dataloader = DataLoader( dataset, batch_size=32, num_workers=4, pin_memory=True, persistent_workers=True # ? 关键！ ) pin_memory=True必须开启：配合non_blocking=True实现零拷贝 DMA 传输，否则.cuda(non_blocking=True)退化为同步操作。避免在训练循环中滥用
empty_cache()：torch.cuda.empty_cache()仅释放未被张量引用的缓存显存，不能释放正在使用的显存。频繁调用反而增加开销，且对子进程泄漏无效。仅在调试或极少数显存碎片严重时（如动态 shape 训练）酌情使用。
? 错误模式复盘（为什么原代码失败）
| 原操作 | 问题本质 |
|---|---|
torch.tensor(...).cuda() 在 __getitem__() 中 |
子进程创建 CUDA 张量 → 独立 CUDA 上下文 → 主进程无法回收其显存 |
del + empty_cache() 在循环内 |
对子进程泄漏无效；且 del 不等于立即释放（Python GC 延迟） |
num_workers=0 临时修复 |
验证了问题根源（多进程 CUDA），但牺牲数据加载吞吐 |
✅ 终极验证与兜底方案
若经上述优化后仍有缓慢增长（罕见），可添加轻量级清理（仅用于诊断/边缘场景）：
import gc # ... 在每个 batch 结束后 del images, labels, output, loss gc.collect() # 强制 Python 垃圾回收 torch.cuda.empty_cache() # 清理主进程空闲缓存
但请牢记：规范的数据流设计（CPU 加载 → 主进程 GPU 迁移）+ persistent_workers=True 已解决 95% 以上的此类泄漏。真正的内存瓶颈通常源于模型层保留历史计算图（如未设 torch.no_grad() 的推理）、梯度累积未清零、或自定义 autograd.Function 中的内存持有，而非数据加载本身。
遵循这一范式，你的训练脚本将获得稳定、可预测且高性能的 GPU 显存行为。
