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

Posted

技术标签:

【中文标题】Tensorflow - 多 GPU 不适用于模型(输入),也不适用于计算梯度【英文标题】:Tensorflow - Multi-GPU doesn’t work for model(inputs) nor when computing the gradients 【发布时间】:2021-09-17 20:25:55 【问题描述】:

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

例如下面这段代码sn-p:

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'])。 按照@Kaveh 的建议添加了这个strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())

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

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

工作示例

您可以在下面的 dogs_vs_cats 数据集上找到一个使用 CNN 训练的工作示例。你不需要手动下载数据集,因为我使用的是 tfds 版本,也不需要训练模型。

笔记本:Working Example.ipynb

保存的模型

HDF5 Saved Format

【问题讨论】:

您可能需要将所有代码放在镜像策略范围内,现在只有模型加载在范围内。 如何确定正在使用一个 GPU? 这个strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce()) 可能会解决您的问题。 你试过在你的定义中列出 gpus 吗?像这样:strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"]). 您能提供一个最小的可重现示例吗?现在,在尝试重现该行为时,我使用了 2 个 GPU。在任何情况下,您可能想查看:tensorflow.org/api_docs/python/tf/distribute/… 【参考方案1】:

对于mirrored_strategy.run() 之外的任何代码,应该在单个 gpu(可能是第一个 gpu,GPU:0)中运行。此外,由于您希望从副本返回渐变,因此还需要 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)

【讨论】:

嗨@Laplace Ricky,非常感谢您的帮助。我刚刚尝试过,这似乎有效!恭喜!这个问题困扰我很久了,感激不尽! 我对您的回答有一些疑问: 1. tf.squeeze(tape.batch_jacobian(outputs, pre_images)) 是做什么的,您为什么不使用tape.gradients?我尝试使用tape.gradient 并且它有效! 2.mirrored_strategy.gather(per_replica_grads,0)中的第二个参数是什么,为什么是零? 由于赏金即将到期并且解决方案有效,我给你奖励。不过,我想进一步研究答案。 回复您的问题,1. 我只是在这里举一个任意的例子,因为我不知道您想要什么渐变。我的示例计算第一个输出(来自 resnet50 的 1000 个输出)相对于输入的梯度。关于tape.gradient,为了清楚起见,我们通常将一个标量传递给它的第一个参数,但如果你认为它给出了你想要的渐变,那就继续吧。 2. 由于分布式数据集将输入数据沿第一维划分为副本,mirrored_strategy.gather 将输出数据沿第一维连接回来。

以上是关于Tensorflow - 多 GPU 不适用于模型(输入),也不适用于计算梯度的主要内容,如果未能解决你的问题,请参考以下文章

具有推理功能的 TensorFlow + Keras 多 GPU 模型

TensorFlow 模型适用于 Python,但不适用于 C++

TensorFlow在使用模型的时候,怎么利用多GPU来提高运算速度

深度学习TensorFlow如何使用多GPU并行模式?

TFLite 转换器:为 keras 模型实现的 RandomStandardNormal,但不适用于纯 TensorFlow 模型

在 LSTM 网络的输入上使用 Masking 时,Keras(TensorFlow 后端)多 GPU 模型(4gpus)失败