如何让TensorFlow的train.py使用所有可用GPU？及资源耗尽问题求助

嘿，我来帮你拆解并解决这几个TensorFlow 1.x训练目标检测时的常见问题：

首先得明确：NAS（神经架构搜索）模型的参数量和计算量比ResNet101大得多，你的双8GB GPU显存确实可能扛不住默认配置。给你几个实用的解决方案：

• 减小批次大小：找到faster_rcnn_nas_coco.config配置文件，修改train_config下的batch_size，从默认值（通常是2）改成1，若仍报错可尝试更小的批次（比如半批次，需微调代码支持）。
• 开启显存动态增长：在train.py开头添加这段代码，让TensorFlow按需分配显存，而非一开始就占满：

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

• 尝试混合精度训练：用TF1.x的tf.contrib.mixed_precision模块，将部分计算转为FP16格式，大幅降低显存占用。需要给优化器套一层LossScaleOptimizer：

from tensorflow.contrib.mixed_precision import LossScaleOptimizer
optimizer = tf.train.MomentumOptimizer(learning_rate=0.001, momentum=0.9)
optimizer = LossScaleOptimizer(optimizer, loss_scale='dynamic')
train_op = optimizer.minimize(loss)

这个提示是官方在推广更健壮的训练会话API——MonitoredTrainingSession比传统的tf.Session好用太多：它会自动处理模型初始化、checkpoint保存/恢复、日志写入、异常终止后的恢复，还天然支持多GPU/分布式训练。

你只需要替换train.py里的会话代码即可：
原来的代码大概是这样：

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for step in range(max_steps):
        sess.run(train_op)

改成：

with tf.train.MonitoredTrainingSession() as sess:
    while not sess.should_stop():
        sess.run(train_op)

这样既消除了提示，也让训练流程更稳定。

TensorFlow 1.x默认会尝试占用所有GPU显存，但有时候需手动配置确保生效：

• 检查环境变量：确保没有设置CUDA_VISIBLE_DEVICES限制GPU使用；若之前设过，要么删除该变量，要么设置为CUDA_VISIBLE_DEVICES=0,1（对应你的两个GPU）。
• 显式指定GPU数量：如果你的train.py支持--num_gpus参数，启动时直接加--num_gpus=2即可，官方目标检测脚本一般都支持这个参数。
• 用分布式策略绑定多GPU：如果脚本无现成参数，就用tf.contrib.distribute.MirroredStrategy实现单机多GPU训练，把模型构建代码放到策略上下文里：

strategy = tf.contrib.distribute.MirroredStrategy()
with strategy.scope():
    # 在这里构建模型、损失、优化器
    rcnn_model = build_faster_rcnn_model()
    loss = rcnn_model.loss()
    optimizer = tf.train.MomentumOptimizer(0.001, 0.9)
    train_op = optimizer.minimize(loss)

# 用MonitoredTrainingSession启动训练
with tf.train.MonitoredTrainingSession() as sess:
    while not sess.should_stop():
        sess.run(train_op)

这样TensorFlow会自动把模型复制到两个GPU上，并行处理数据，充分利用硬件资源。

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