Tensorflow 中的多 GPU 训练在 Nans 中得到结果

Question

我正在尝试使用多个 GPU 进行训练，但是几步之后损失始终为 Nan。如果我使用单个 GPU，那就没问题了。下面显示了一个虚拟脚本，经过几个步骤后会生成 nan。下面的代码显示了我放置在顶部的用于张量流构建信息/GPU 信息的打印语句的输出。

tensorflow版本是2.18.0

import tensorflow as tf
import numpy as np

print(tf.sysconfig.get_build_info())

# Check if GPUs are available
gpus = tf.config.list_physical_devices('GPU')
print(gpus)
if gpus:
    print(f"Number of GPUs available: {len(gpus)}")
else:
    print("No GPUs found. Training will proceed on CPU.")

# Define the strategy for multi-GPU training
strategy = tf.distribute.MirroredStrategy()

# Dummy dataset
def create_dummy_dataset(samples=50000):
    # Generate dummy input data 
    X = np.random.random((samples, 20)).astype(np.float32)
    # Generate dummy labels (binary classification)
    y = np.random.randint(0, 2, (samples, 1)).astype(np.float32)
    return tf.data.Dataset.from_tensor_slices((X, y)).shuffle(samples).batch(32)

dataset = create_dummy_dataset()

# Define the model inside the strategy scope
with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(64, activation='relu', input_shape=(20,)),
        tf.keras.layers.Dense(32, activation='relu'),
        tf.keras.layers.Dense(1, activation='sigmoid')
    ])
    model.compile(optimizer='adam',
                  loss='binary_crossentropy',
                  metrics=['accuracy'])

# Train the model
model.fit(dataset, epochs=10)

OrderedDict(
{'cpu_compiler': '/usr/lib/llvm-18/bin/clang', 'cuda_compute_capabilities': ['sm_60', 'sm_70', 'sm_80', 'sm_89', 'compute_90'], 'cuda_version': '12.5.1', 'cudnn_version': '9', 'is_cuda_build': True, 'is_rocm_build': False, 'is_tensorrt_build': False}
)

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU')] 
Number of GPUs available: 2 

Answer 1

我无法发表评论，所以我将此作为答案。

如果我使用单个 GPU，那就没问题了。下面显示了一个虚拟脚本，经过几个步骤后会生成 nan。

我认为这可能是由于您的批量大小所致；尝试增加它，因为它会让你的损失更加稳定。另外，您用于单 GPU 训练的批量大小是多少？

https://www.tensorflow.org/tutorials/distribute/keras#set_up_the_input_pipeline

如果您检查上面的链接，您可以看到下面的代码行。

BATCH_SIZE = BATCH_SIZE_PER_REPLICA * Strategy.num_replicas_in_sync

希望这有帮助。