ML模型训练时系统内存随Epoch持续增长的问题解决求助

核心问题定位 #

当num_workers>0时内存持续增长，说明数据加载子进程存在内存泄漏，和模型训练逻辑无关。以下是针对性的解决手段：

1. 修正DataLoader的worker初始化逻辑 #

• 禁止在__getitem__方法中加载大文件或创建全局变量：子进程会保留这些对象的引用，无法随批次自动释放。比如提前在数据集的__init__中加载预训练嵌入等大资源，共享给所有worker，不要在每个样本加载时重复初始化。
• 自定义worker_init_fn清理子进程状态：给DataLoader传入初始化函数，重置全局状态并触发垃圾回收，示例代码：

  def worker_init_fn(worker_id):
      torch.seed()
      import gc
      gc.collect()
  
  DataLoader(dataset, num_workers=4, worker_init_fn=worker_init_fn)
  

2. 合理配置worker数量与持久化 #

• 匹配CPU核心数设置worker数量：ml.g4.xlarge有4个vCPU，64个worker远超硬件承载能力，会导致进程切换开销和内存暴涨。建议设为num_workers=4（等于核心数）或num_workers=8（2倍核心数），平衡训练速度与内存占用。
• 启用persistent_workers=True：让worker在Epoch间保持存活，避免重复初始化带来的内存累积，但必须配合worker_init_fn清理每个Epoch后的状态，否则会加剧泄漏。

3. 主动触发内存回收 #

• 在Epoch结束时强制回收内存：在PyTorch Lightning的on_epoch_end钩子中添加垃圾回收逻辑，同时清理GPU显存碎片：

  def on_epoch_end(self):
      import gc
      gc.collect()
      torch.cuda.empty_cache()
  

4. 优化数据预处理流程 #

• 避免在预处理中生成大临时对象：比如不要在__getitem__中创建未及时释放的大numpy数组，尽量用PyTorch张量直接处理，减少内存拷贝。
• 开启pin_memory=True：如果数据需要从CPU传输到GPU，开启该参数能优化内存传输效率，减少内存碎片，但要确保预处理后的张量无冗余引用。

5. SageMaker环境专属配置 #

• 启用swap内存缓解内存压力：执行以下命令创建16G交换分区（按需调整大小）：

  sudo fallocate -l 16G /swapfile
  sudo chmod 600 /swapfile
  sudo mkswap /swapfile
  sudo swapon /swapfile
  

• 限制日志缓存：设置环境变量减少日志组件的内存占用：

  import os
  os.environ['SM_LOG_LEVEL'] = 'ERROR'
  os.environ['PYTHONUNBUFFERED'] = '1'
  

内容的提问来源于stack exchange，提问作者Manuel Ravasqueira