TensorFlow-Keras 2.X 多GPU训练正常但预测仅单GPU运行解决方法

单卡运行预测的核心原因是：预测阶段没有将模型加载、推理流程纳入tf.distribute.MirroredStrategy的作用域，TensorFlow默认不会自动调度多卡做分布式推理。

具体配置步骤 #

• 预测阶段和训练阶段保持一致，先初始化MirroredStrategy分布式策略
• 所有模型相关操作（包括模型构建、训练好的权重/整模加载）必须放在策略的scope()作用域内完成，禁止在作用域外加载模型再传入做推理
• 调整预测输入的总batch size为单卡最优batch size * 4（你的GPU总数），否则单卡分到的计算量太小，会出现GPU利用率低、看起来只有单卡运行的情况
• 优先使用model.predict()配合tf.data输入流做推理，手动循环调用predict_on_batch()不会自动触发分布式调度逻辑

可直接复用的代码示例 #

import tensorflow as tf

# 初始化分布式策略，和训练阶段配置保持一致
mirrored_strategy = tf.distribute.MirroredStrategy()

# 模型加载/构建必须放在策略作用域内完成
with mirrored_strategy.scope():
    # 二选一：训练时保存完整模型的场景，直接加载整模
    # model = tf.keras.models.load_model(your_trained_model_path)
    # 二选一：用ModelCheckpoint保存权重的场景，先构建模型结构再加载权重
    model = MyModel()
    model.load_weights(your_checkpoint_weight_path)

# 构造测试输入流，总batch size设为单卡最优batch * 4
# 例：单卡推理最优batch为16时，4卡总batch设为64，框架会自动均匀切分到每张卡
# 开启prefetch提前加载数据，避免GPU等数据拖慢速度
testData = testData.batch(64).prefetch(tf.data.AUTOTUNE)

# 直接调用predict即可，框架自动将数据分片调度到4张GPU并行计算
# 返回结果会自动聚合，和单卡推理的输出格式完全一致，无需手动拼接多卡结果
pred_result = model.predict(testData, verbose=1)

常见异常排查 #

• 程序启动后先执行print(tf.config.list_physical_devices('GPU'))，确认可以识别到全部4张RTX 3090，检查是否存在CUDA_VISIBLE_DEVICES环境变量限制了可见GPU数量
• 模型加载阶段只有主卡出现显存占用是正常现象，等predict正式执行计算时，4张卡的显存占用、GPU利用率都会同步上涨，不要在模型加载阶段就判定多卡调度失败
• 不要在代码中手动调用tf.config.set_visible_devices()限制可用GPU数量

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