在Object Detection API中同时训练评估ssd_mobile_v1_coco遇GPU内存不足

我之前也碰到过一模一样的情况！当你在两个终端同时跑训练和评估脚本时，两个TensorFlow进程会疯狂抢占GPU显存，直接把显存榨干导致评估脚本启动失败。给你几个亲测有效的解决办法：

1. 限制每个进程的GPU显存占用 #

你可以给训练和评估进程分别设置显存上限，让它们各用一部分显存，避免互相挤兑。

• 启动脚本时直接加参数（适用于TF 2.x）：

  # 训练脚本，限制用4GB显存
  python train.py --memory_limit 4096 --其他训练参数
  # 评估脚本，限制用3GB显存
  python eval.py --memory_limit 3072 --其他评估参数
  

• 或者在代码开头手动设置显存限制：

  import tensorflow as tf
  gpus = tf.config.experimental.list_physical_devices('GPU')
  if gpus:
      try:
          # 给当前进程分配4GB显存
          tf.config.experimental.set_virtual_device_configuration(
              gpus[0],
              [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4096)])
      except RuntimeError as e:
          print(e)
  

  数值根据你的GPU总显存调整，比如你的GPU是8GB，分给训练4GB、评估3GB，留1GB给系统就行。

2. 交替运行训练和评估，不同时占用GPU #

如果显存实在紧张，不如让训练和评估轮流跑。写个简单的shell脚本就能实现：

# 先训练5000步
python train.py --num_steps 5000 --checkpoint_dir ./train_logs --其他参数

# 运行评估
python eval.py --checkpoint_dir ./train_logs --eval_dir ./eval_logs --其他参数

# 继续训练到10000步
python train.py --num_steps 10000 --checkpoint_dir ./train_logs --其他参数

# 再次评估
python eval.py --checkpoint_dir ./train_logs --eval_dir ./eval_logs --其他参数

这样同一时间只有一个进程占用GPU，显存压力瞬间就小了。

3. 调小训练和评估的batch size #

batch size是显存占用的关键因素之一。把训练的--batch_size从默认值（比如32）调到8或者4，评估的batch size也调小到2或者4，这样每个进程需要的显存会大幅减少，就能同时跑起来了。

4. 多GPU环境下分配不同GPU #

如果你有多块GPU，直接把训练和评估分到不同的GPU上，彻底避免显存冲突：

# 训练脚本用第0块GPU
CUDA_VISIBLE_DEVICES=0 python train.py --其他参数

# 评估脚本用第1块GPU
CUDA_VISIBLE_DEVICES=1 python eval.py --其他参数

这些方法我都试过，根据自己的GPU情况选一个就行。比如我只有一块6GB的GPU，用方法1+3就完美解决了问题。

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