如何设置max_split_size_mb解决Colab Pro+上CUDA内存不足问题？

设置max_split_size_mb的两种有效方式 #

• 环境变量前置设置（Colab中优先用这个）
  在导入PyTorch之前，通过环境变量配置内存分配规则，确保设置生效：

  import os
  os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
  # 之后再导入PyTorch
  import torch
  

  512是初始建议值，可以根据实际情况调整（比如256、1024），核心是限制单次内存分配的块大小，减少碎片堆积。

• 运行时动态调整（PyTorch 1.10及以上版本支持）
  如果已经导入了PyTorch，也可以直接修改分配器设置：

  import torch
  torch.cuda.memory._set_allocator_settings('max_split_size_mb:512')
  

针对Colab Pro+场景的额外内存优化建议 #

• 定期清理GPU缓存：在每个训练epoch结束后，调用torch.cuda.empty_cache()释放未使用的显存，但不要在batch循环内频繁调用，避免拖慢训练速度。
• 及时释放无用张量：对于中间计算结果、验证集临时数据，用完后用del tensor删除，再配合torch.cuda.empty_cache()释放显存。
• 开启混合精度训练：用torch.cuda.amp模块实现混合精度，能大幅降低内存占用，示例代码：

  scaler = torch.cuda.amp.GradScaler()
  for epoch in range(num_epochs):
      for batch in dataloader:
          inputs, labels = batch
          inputs, labels = inputs.cuda(), labels.cuda()
          with torch.autocast(device_type='cuda'):
              outputs = model(inputs)
              loss = criterion(outputs, labels)
          scaler.scale(loss).backward()
          scaler.step(optimizer)
          scaler.update()
  

• 确认h5py数据加载逻辑：避免将整个数据集预加载到GPU，应该从h5py读取到CPU张量，再转移到GPU，减少显存占用。

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