使用tf.contrib.distribute.MirroredStrategy未调用多GPU训练问题咨询

针对你遇到的MirroredStrategy检测到GPU但未实际使用的问题，这在TensorFlow 1.x早期版本的分布式策略使用中是一个常见配置问题，核心原因是模型和优化器的创建时机没有在策略的作用域内。以下是具体的解决方案和解释：

问题根源 #

你当前的代码先创建了Keras模型、优化器并完成编译，之后才初始化MirroredStrategy并转换为Estimator。这种顺序会导致策略无法捕获模型的分布式训练逻辑——TensorFlow 1.x的分布式策略需要在其作用域内定义所有训练相关组件，才能正确将计算分配到各个GPU上。

修正后的代码示例 #

import tensorflow as tf
from tensorflow.keras import models, Input

LEARNING_RATE = 0.001
NUM_GPUS = 2

# 1. 先初始化分布式策略
strategy = tf.contrib.distribute.MirroredStrategy(num_gpus=NUM_GPUS)

# 2. 在策略的作用域内创建模型、优化器并编译
with strategy.scope():
    # 替换为你的实际输入形状与模型结构
    input = Input(shape=(28, 28, 1))
    x = tf.keras.layers.Flatten()(input)
    x = tf.keras.layers.Dense(128, activation='relu')(x)
    y_output = tf.keras.layers.Dense(10, activation='softmax')(x)
    
    model = models.Model(inputs=input, outputs=y_output)
    optimizer = tf.train.AdamOptimizer(LEARNING_RATE)
    model.compile(loss=lossFunc, optimizer=optimizer)

# 3. 转换为Estimator并启动训练
config = tf.estimator.RunConfig(train_distribute=strategy)
estimator = tf.keras.estimator.model_to_estimator(model, config=config)

额外注意事项 #

• 版本兼容性检查：确保TensorFlow 1.12与CUDA、cuDNN版本匹配（TF1.12要求CUDA 9.0，cuDNN 7.1.x），版本不匹配可能导致GPU无法被正确调用。
• 环境变量验证：检查是否设置了CUDA_VISIBLE_DEVICES环境变量，避免意外屏蔽GPU设备；如果手动指定该变量，确保包含你要使用的GPU编号。
• 自定义组件兼容性：如果你的lossFunc是自定义函数，建议确保它能被分布式策略正确序列化，部分复杂逻辑可能需要调整以适配早期分布式实现。

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