如何强制TensorFlow使用所有可用GPU？8GPU集群单卡运行问题解决

Hey there, let's tackle that frustrating issue where your TensorFlow model is only using one GPU in your 8-GPU cluster—you confirmed it via nvidia-smi, so let's get all those GPUs pulling their weight!

1. Use TensorFlow's MirroredStrategy (Best for Single-Node Multi-GPU) #

TensorFlow doesn't automatically distribute training across GPUs by default, but the MirroredStrategy is the easiest way to fix this for a single machine with multiple GPUs. It replicates your model on every GPU and splits your training data into chunks for each GPU to process.

Here's how to modify your code:

import tensorflow as tf
from tensorflow.keras import layers, models, optimizers

# Initialize the strategy to use all available GPUs
strategy = tf.distribute.MirroredStrategy()

# All model building and compilation MUST happen inside this scope
with strategy.scope():
    inputs = layers.Input((IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS))
    # ... paste your full model architecture here (conv layers, etc.) ...
    outputs = layers.Conv2D(1, (1, 1), activation='sigmoid')(c9)
    model = models.Model(inputs=[inputs], outputs=[outputs])
    
    sgd = optimizers.SGD(lr=0.03, decay=1e-6, momentum=0.9, nesterov=True)
    # Make sure your metrics are properly defined (e.g., mean_iou should be tf.keras.metrics.MeanIoU)
    model.compile(optimizer=sgd, loss='binary_crossentropy', metrics=[tf.keras.metrics.MeanIoU(num_classes=2)])

# When you train, the strategy handles splitting data across GPUs
model.fit(train_data, epochs=10, validation_data=val_data)

2. First, Confirm TensorFlow Sees All 8 GPUs #

Before diving into distribution, run this quick check to make sure TensorFlow detects all your GPUs:

print("Number of GPUs detected by TensorFlow: ", len(tf.config.list_physical_devices('GPU')))

If this returns less than 8, double-check your TensorFlow GPU installation and NVIDIA driver setup—you might have a setup issue preventing TensorFlow from accessing all cards.

3. Adjust Your Batch Size #

When using multiple GPUs, scale your batch size up proportionally. For example, if you were using a batch size of 32 on one GPU, try 32 * 8 = 256. This ensures each GPU gets a similar workload as your original single-GPU setup, keeping training efficient.

4. Avoid Common Mistakes #

• Don't build the model outside the strategy scope: Every part of model creation, compilation, and even dataset preparation (if using tf.data) needs to live inside strategy.scope() to ensure proper distribution.
• Remove hard-coded GPU locks: If your code has lines like tf.config.set_visible_devices(gpus[0], 'GPU'), delete them—this forces TensorFlow to only use the first GPU.
• Use tf.data.Dataset for inputs: This pipeline is optimized for distributed training and automatically splits data across GPUs, which is more reliable than using numpy arrays directly.

5. Verify Usage During Training #

Once you start training, run nvidia-smi in a separate terminal. You should see all 8 GPUs showing memory usage and GPU utilization—no more idle cards!

That should get your cluster running at full capacity. If you hit any roadblocks with the code adjustments, feel free to troubleshoot further.

内容的提问来源于stack exchange，提问作者Jonathan