TensorFlow多GPU下调用model(inputs)及计算梯度仅单卡运行问题

使用多GPU对模型执行推理（例如调用model(inputs)方法）并计算梯度时，机器仅使用单张GPU，其余显卡处于完全空闲状态。
示例代码片段如下：

import tensorflow as tf
import numpy as np
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

# 构造tf-data
path_filename_records = 'your_path_to_records'
bs = 128

dataset = tf.data.TFRecordDataset(path_filename_records)
dataset = (dataset
           .map(parse_record, num_parallel_calls=tf.data.experimental.AUTOTUNE)
           .batch(bs)
           .prefetch(tf.data.experimental.AUTOTUNE)
          )

# 加载使用MirroredStrategy训练的模型
path_to_resnet = 'your_path_to_resnet'
mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
    resnet50 = tf.keras.models.load_model(path_to_resnet)

for pre_images, true_label in dataset:
    with tf.GradientTape() as tape:
       tape.watch(pre_images)
       outputs = resnet50(pre_images)
       grads = tape.gradient(outputs, pre_images)

运行上述代码时仅有单张GPU被使用，可通过nvidia-smi命令监控GPU运行状态。目前无法确认model(inputs)和tape.gradient本身是否不支持多GPU运行，若确实存在该限制会带来较大性能问题：当存在大规模数据集需要计算相对于输入的梯度（例如可解释性相关需求）时，单卡运行可能需要数天时间。另外尝试过使用model.predict()方法，但该方法无法与tf.GradientTape搭配使用。

已尝试但未生效的方案 #

• 将所有代码放在MirroredStrategy作用域内
• 更换不同型号GPU：已测试A100、A6000、RTX5000，同时调整过显卡数量和批次大小
• 显式指定GPU列表，例如strategy = tf.distribute.MirroredStrategy(['/gpu:0', '/gpu:1'])
• 配置跨设备操作参数：strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())

仅单卡运行的验证方式 #

在终端使用watch -n 1 nvidia-smi命令实时监控GPU状态，可见仅有一张GPU使用率为100%，其余显卡使用率均为0%。

可复现示例 #

提供基于猫狗数据集训练的CNN可运行示例，使用tfds版本数据集，无需手动下载数据集，也不需要自行训练模型。

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