tf.estimator高层API与tf.slim兼容性问题：allow_growth参数失效致GPU内存耗尽

我之前也碰到过一模一样的坑——用tf.estimator搭NASNet（基于tf.slim实现）时，明明设置了GPU内存增长模式，结果显卡还是被瞬间占满。核心原因是tf.estimator默认会使用内置的会话配置，不会自动继承你在全局设置的GPU选项，必须显式把自定义配置传给Estimator才行。

下面是针对你场景的具体解决方案：

1. 构建自定义GPU会话配置 #

先创建tf.ConfigProto实例，开启内存按需分配的模式，如果需要限制内存占比也可以额外配置：

import tensorflow as tf

# 配置GPU内存策略
gpu_config = tf.ConfigProto()
# 开启allow_growth，让TensorFlow按需分配GPU内存
gpu_config.gpu_options.allow_growth = True
# 可选：限制GPU内存使用率（比如只占用70%），不需要可以注释掉
# gpu_config.gpu_options.per_process_gpu_memory_fraction = 0.7

2. 创建Estimator的RunConfig #

把上面的GPU配置传入tf.estimator.RunConfig，同时可以指定模型保存路径、 checkpoint保存频率等参数：

run_config = tf.estimator.RunConfig(
    session_config=gpu_config,
    model_dir="./nasnet_checkpoints",  # 替换成你的模型保存路径
    save_checkpoints_steps=1000
)

3. 初始化Estimator时传入RunConfig #

在创建你的Estimator实例时，一定要把run_config作为config参数传进去，替换掉默认配置：

# 用你自己定义的model_fn和参数字典初始化Estimator
estimator = tf.estimator.Estimator(
    model_fn=model_fn,
    params=your_params,  # 替换成你的模型参数字典
    config=run_config
)

额外提醒 #

• 如果是TensorFlow 2.x环境（虽然你用的是tf.estimator和slim，大概率是1.x，但还是提一下），要对应使用tf.compat.v1.ConfigProto和tf.compat.v1.estimator.RunConfig。
• 开启allow_growth后，第一次运行模型时可能还是会占用较多内存，但这是模型初始化的正常现象，后续会按需动态调整，不会一直占满整个GPU。

这样配置后，Estimator就会遵循你设置的GPU内存规则，不会再一启动就吃光显卡内存了。

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