AWS EC2 g4dn.xlarge实例运行TensorFlow GPU Python代码自动被kill问题求解

完全可以在当前g4dn.xlarge实例规格下解决该问题，核心思路是限制TensorFlow的GPU显存占用、优化内存释放逻辑、利用闲置的CPU内存资源分摊GPU压力，具体操作如下：

1 先确认问题根因 #

• 运行代码时同时开终端执行nvidia-smi -l 1实时监控显存占用，配合dmesg | grep -E 'Killed process|oom-killer'查看系统日志是否有OOM杀死进程的记录，可100%验证显存溢出的猜测。

2 TensorFlow 2.4 专属显存优化配置 #

你使用的TF2.4版本默认会预占几乎全部GPU显存，将以下代码加到所有导入TensorFlow逻辑的最前面，可大幅降低无效显存占用：

import tensorflow as tf
# 限制TensorFlow仅按需分配显存
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
    # 也可以直接硬限制最大显存占用，比如限制为14GB，留2GB余量防止溢出
    # tf.config.experimental.set_virtual_device_configuration(
    #     gpus[0],
    #     [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=14336)])
  except RuntimeError as e:
    print(e)

3 优化图像处理逻辑的显存释放 #

你采用的递归处理方式容易出现前一组处理完的张量没有被及时回收，递归调用栈也会持有变量引用导致GC无法释放，可做如下调整：

• 把递归调用改成循环迭代处理每组图像，避免调用栈累积变量
• 每组处理完成后显式调用tf.keras.backend.clear_session()清空Keras会话里的所有张量，再配合import gc; gc.collect()强制回收Python内存和显存
• 单组图像处理的batch size下调30%-50%，哪怕多跑几个batch，也不要一次性往GPU塞太多数据，小batch跑的总速度不会比大batch慢多少，还能大幅降低峰值显存占用

4 利用闲置CPU内存分摊GPU压力 #

你当前CPU和系统内存有90%空闲，可以把部分预处理逻辑放到CPU执行，减少GPU的显存占用：

• 图像的读取、缩放、归一化等预处理步骤全部放在CPU侧完成，不要放到TF的GPU张量里做
• 如果用的是tf.data做数据加载，加上.prefetch(tf.data.AUTOTUNE)的同时指定tf.data.experimental.ignore_errors()，并且把数据缓存到系统内存而不是GPU显存：dataset = dataset.cache("/tmp/cache")

5 可选AWS专属优化（无需升配） #

如果以上逻辑优化完还有偶尔的OOM，可以用AWS自带的服务特性兜底，不需要换实例：

• 给EC2实例加一块交换分区（Swap），用闲置的EBS存储当虚拟内存，哪怕显存溢出也会先换到虚拟内存，不会直接被杀死，g4dn.xlarge默认没有Swap，执行以下命令配置16GB Swap即可：

  sudo fallocate -l 16G /swapfile
  sudo chmod 600 /swapfile
  sudo mkswap /swapfile
  sudo swapon /swapfile
  # 开机自动挂载Swap
  echo '/swapfile none swap sw 0 0' | sudo tee -a /etc/fstab
  

• 如果你处理的任务可以拆分，用AWS Batch把任务拆成多个小任务，每个小任务处理一组图像，跑完就释放进程，完全避免长时间运行的进程显存泄漏问题，只需要用当前的g4dn.xlarge作为计算节点即可，不会产生额外成本。

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