Google Colab训练CNN时首个epoch后因内存分配报错自动重启

同一份代码隔天跑触发内存报错重启，核心原因是Colab免费实例的硬件配额是动态分配的，前一天运行时你大概率分到了16G以上显存的高内存GPU实例，本次运行分到的是低配置标准实例，加上训练流程里的内存峰值刚好卡在实例的OOM阈值线上，第一个epoch跑完进入验证、指标统计环节时内存占用触顶，系统会直接强制终止内核重启，很多时候不会抛出完整的内存报错栈。

可按以下顺序排查修复：

• 先确认当前实例的硬件配额：
  运行命令查看系统总内存：!cat /proc/meminfo | grep MemTotal
  运行命令查看GPU显存：!nvidia-smi
  如果查到系统内存低于12G、GPU显存低于16G，直接在菜单栏选「修改-笔记本设置」重新连接实例，多尝试2-3次就能分到配置足够的实例，免费实例资源分配本身存在随机性，同代码不同运行时间分到的硬件配置差异很大是正常情况。
• 在训练循环中加入显式内存回收逻辑，避免张量残留占内存：
  在每个epoch执行完的位置加入缓存清理代码，根据你用的框架选对应版本：
  PyTorch框架用：

  import gc
  import torch
  gc.collect()
  torch.cuda.empty_cache()
  

  TensorFlow/Keras框架用：

  import gc
  import tensorflow as tf
  gc.collect()
  tf.keras.backend.clear_session()
  

  注意不要把所有批次的中间特征图、训练日志全存在全局列表里，日志只保留每个epoch的聚合指标即可，中间张量用完如果不需要回溯就及时del掉。
• 压低训练流程的内存峰值：
  第一个epoch结束时的内存峰值一般出现在验证环节，把验证阶段的batch size调小到训练batch size的1/2；如果之前是把整个数据集一次性加载到内存，改成用DataLoader做懒加载，不要全量读入所有训练数据；训练时关闭不必要的实时可视化、动态样本增强缓存，避免额外占用内存。
• 清理实例内的闲置进程：
  跑训练前先执行命令杀掉占内存的后台闲置进程：

  !pkill -f chromium
  !fuser -k 8888/tcp
  

  不要在跑训练的同时执行大文件云盘同步、其他数据处理脚本，免费实例的OOM触发阈值非常严格，内存占用超过90%持续数秒就会直接杀内核，不会预留报错时间。

如果你是Colab Pro用户，可以直接在运行时类型里勾选「高内存」选项，能直接规避90%以上的这类随机OOM问题。

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