TensorFlow与PyTorch混合流水线重置内存统计时CUDA OOM问题排查

你遇到的核心问题是TensorFlow的GPU内存配置代码执行时机太晚。TF在你第一次调用tf.config.list_logical_devices("GPU")时就已经完成了GPU初始化，此时再设置set_memory_growth或set_logical_device_configuration根本不会生效，TF还是会默认占满GPU内存，导致PyTorch初始化时没有剩余内存可用。

必须把TF的GPU内存配置代码放在所有TF设备查询、模型加载、张量操作之前，确保TF初始化时就遵循你设定的内存规则。

方案1：启用内存按需增长（推荐） #

import tensorflow as tf
import torch

# 第一步：先配置TF的GPU内存增长，必须在任何TF设备操作前执行
for gpu in tf.config.list_physical_devices("GPU"):
    tf.config.experimental.set_memory_growth(gpu, True)

# 之后再执行TF的内存统计重置
for x in tf.config.list_logical_devices("GPU"):
    tf.config.experimental.reset_memory_stats(x.name[x.name.index(":") + 1:])

# 最后初始化PyTorch并重置其内存统计
torch.cuda.init()
for x in range(torch.cuda.device_count()):
    torch.cuda.reset_peak_memory_stats(x)

方案2：固定TF的GPU内存上限 #

如果需要严格限制TF的内存使用量，同样要把配置放在最前面：

import tensorflow as tf
import torch

# 第一步：先设置TF的GPU内存限制
for gpu in tf.config.list_physical_devices("GPU"):
    tf.config.set_logical_device_configuration(
        gpu,
        [tf.config.LogicalDeviceConfiguration(memory_limit=1024)]  # 这里设为1GB，按需调整数值
    )

# 后续操作不变
for x in tf.config.list_logical_devices("GPU"):
    tf.config.experimental.reset_memory_stats(x.name[x.name.index(":") + 1:])

torch.cuda.init()
for x in range(torch.cuda.device_count()):
    torch.cuda.reset_peak_memory_stats(x)

• 如果你在脚本中还有其他TF相关的导入或操作（比如加载模型、定义张量），也要确保内存配置代码在这些操作之前执行。
• 单独运行TF或PyTorch的代码没问题，是因为单个框架不会抢占对方的内存空间，但混合使用时必须严格控制初始化顺序和内存策略。

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