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

Question

当使用多个 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 来说是不可能的.
到目前为止我已经尝试过但没有奏效

将所有代码放在镜像策略范围内。

使用不同的 GPU:我尝试过 A100、A6000 和 RTX5000。还更改了图形卡的数量并改变了批量大小。

指定 GPU 列表，例如 strategy = tf.distribute.MirroredStrategy(['/gpu:0', '/gpu:1']) .

添加此 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 版本那样手动下载数据集，也不需要训练模型。
笔记本 :Working Example.ipynb
已保存模型 :

HDF5

Saved Format

Answer 1

是假设 对于在 GPU:0 之外的任何代码，在单个 GPU(可能是第一个 GPU， mirrored_strategy.run() )中运行.此外，由于您希望从副本返回梯度，mirrored_strategy.gather()也需要。
除此之外，必须使用 mirrored_strategy.experimental_distribute_dataset 创建分布式数据集。 .分布式数据集尝试在副本之间均匀分布单批数据。以下是关于这些要点的示例。model.fit() , model.predict() ，等等...以分布式方式自动运行，因为他们已经为您处理了上面提到的所有内容。
示例代码:

mirrored_strategy = tf.distribute.MirroredStrategy()
print(f'using distribution strategy\nnumber of gpus:{mirrored_strategy.num_replicas_in_sync}')

dataset=tf.data.Dataset.from_tensor_slices(np.random.rand(64,224,224,3)).batch(8)

#create distributed dataset
ds = mirrored_strategy.experimental_distribute_dataset(dataset)

#make variables mirrored
with mirrored_strategy.scope():
  resnet50=tf.keras.applications.resnet50.ResNet50()

def step_fn(pre_images):
  with tf.GradientTape(watch_accessed_variables=False) as tape:
       tape.watch(pre_images)
       outputs = resnet50(pre_images)[:,0:1]
  return tf.squeeze(tape.batch_jacobian(outputs, pre_images))

#define distributed step function using strategy.run and strategy.gather
@tf.function
def distributed_step_fn(pre_images):
  per_replica_grads = mirrored_strategy.run(step_fn, args=(pre_images,))
  return mirrored_strategy.gather(per_replica_grads,0)

#loop over distributed dataset with distributed_step_fn
for result in map(distributed_step_fn,ds):
  print(result.numpy().shape)