Tensorflow-多GPU不适用于模型（输入）或计算梯度

当使用多个 GPU 对模型进行推理（例如调用方法：model(inputs)）并计算其梯度时，机器只使用一个 GPU，其余空闲。

例如在下面的代码片段中：

import tensorflow as tf
import numpy as np
import os

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

# Make the 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)
          )

# Load model trained using 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。但如果是这样，那就是一个大问题，因为如果您有一个大数据集并且需要计算关于输入的梯度（例如可解释性海豚），使用一个 GPU 可能需要几天时间。我尝试过的另一件事是使用，model.predict()但这对于tf.GradientTape.

到目前为止我已经尝试过但没有奏效

1. 将所有代码放在镜像策略范围内。
2. 使用不同的 GPU：我尝试过 A100、A6000 和 RTX5000。还更改了图形卡的数量并改变了批量大小。
3. 指定 GPU 列表，例如，strategy = tf.distribute.MirroredStrategy(['/gpu:0', '/gpu:1'])。
4. strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())按照@Kaveh 的建议添加了此内容。

我怎么知道只有一个 GPU 在工作？

我watch -n 1 nvidia-smi在终端中使用了该命令并观察到只有一个 GPU 处于 100%，其余的处于 0%。

工作示例

您可以在下面的 dog_vs_cats 数据集上找到一个使用 CNN 训练的工作示例。您不需要像我使用 tfds 版本那样手动下载数据集，也不需要训练模型。

笔记本：工作示例.ipynb

保存的模型：

• HDF5
• 保存格式