当使用多个GPU对模型进行推理(例如调用方法:model(inputs))并计算其梯度时,机器只使用一个GPU,其余的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)
and 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%。
工作示例
您可以在下面找到一个在 dogs_vs_cats 数据集上训练的 CNN 的工作示例。您不需要像我使用 tfds 版本一样手动下载数据集,也不需要训练模型。
Notebook: 工作示例.ipynb https://drive.google.com/file/d/1CoYVckmEQXp2Wf_PRlGtnrlGnF8smxvt/view?usp=sharing
已保存模型:
- HDF5 https://drive.google.com/file/d/1Y0-fQytVsnHPs8JL6kKJr3tEL0mKNtjJ/view?usp=sharing
- 保存格式 https://drive.google.com/file/d/19oSIaUTtEy1q6rlDj8GzuWIvwAq_7XMi/view?usp=sharing
It is supposed在单个 GPU 中运行(可能是第一个 GPU,GPU:0
)对于任何超出范围的代码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)
本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)