Multi-GPU training basics in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Multi-GPU training basics

Problem:You have a neural network model training on a single GPU. Training is slow and you want to speed it up by using multiple GPUs.

Current Metrics:Training time per epoch: 120 seconds, Validation accuracy: 82%, Training accuracy: 90%

Issue:Training is slow due to single GPU usage. Model training time is long, but accuracy is reasonable.

Your Task

Use TensorFlow's multi-GPU strategy to reduce training time per epoch by at least 40% while maintaining validation accuracy above 80%.

Do not change the model architecture.

Keep batch size reasonable (e.g., 64 or 128).

Use TensorFlow 2.x APIs only.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

TensorFlow

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

# Detect available GPUs
physical_gpus = tf.config.list_physical_devices('GPU')
print(f"GPUs available: {len(physical_gpus)}")

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

with strategy.scope():
    # Define a simple CNN model
    model = models.Sequential([
        layers.Conv2D(32, (3,3), activation='relu', input_shape=(28,28,1)),
        layers.MaxPooling2D((2,2)),
        layers.Conv2D(64, (3,3), activation='relu'),
        layers.MaxPooling2D((2,2)),
        layers.Flatten(),
        layers.Dense(64, activation='relu'),
        layers.Dense(10, activation='softmax')
    ])

    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

# Load MNIST data
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

# Normalize and reshape
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

batch_size = 128

# Measure training time
start_time = time.time()

history = model.fit(x_train, y_train,
                    epochs=5,
                    batch_size=batch_size,
                    validation_data=(x_test, y_test))

end_time = time.time()
training_time_per_epoch = (end_time - start_time) / 5

# Evaluate final accuracy
val_loss, val_accuracy = model.evaluate(x_test, y_test, verbose=0)

print(f"Training time per epoch: {training_time_per_epoch:.2f} seconds")
print(f"Validation accuracy: {val_accuracy*100:.2f}%")

Added tf.distribute.MirroredStrategy to enable multi-GPU training.

Wrapped model creation and compilation inside strategy.scope() to distribute the model.

Increased batch size to 128 to better utilize multiple GPUs.

Measured training time per epoch to compare speed improvements.

Results Interpretation

Before Multi-GPU: Training time per epoch: 120 seconds, Validation accuracy: 82%, Training accuracy: 90%

After Multi-GPU: Training time per epoch: 70 seconds, Validation accuracy: 81.5%, Training accuracy: 89%

Using multiple GPUs with TensorFlow's MirroredStrategy can significantly reduce training time while maintaining similar accuracy. This shows how distributing work across GPUs speeds up training without hurting model performance.

Bonus Experiment

Try increasing the batch size further to 256 or 512 and observe how training time and accuracy change.

💡 Hint

Larger batch sizes can improve GPU utilization but may affect model convergence. Adjust learning rate if accuracy drops.

Practice

1. What is the main purpose of using tf.distribute.MirroredStrategy() in TensorFlow?

easy

A. To train a model on multiple GPUs simultaneously

B. To reduce the size of the dataset

C. To convert a model to run on CPU only

D. To save the model in a special format

Solution

1. Step 1: Understand the role of MirroredStrategy

   tf.distribute.MirroredStrategy() is designed to distribute training across multiple GPUs to speed up the process.

2. Step 2: Compare options with the purpose

   Options A, B, and C do not relate to multi-GPU training. Only training a model on multiple GPUs simultaneously correctly states the purpose.

3. Final Answer:

   To train a model on multiple GPUs simultaneously -> Option A

4. Quick Check:

   Multi-GPU training = MirroredStrategy usage [OK]

Hint: MirroredStrategy means training on many GPUs at once [OK]

Common Mistakes:

• Thinking it reduces dataset size
• Confusing with CPU-only training
• Assuming it changes model saving format

2. Which of the following is the correct way to start using multi-GPU training in TensorFlow?

easy

A. model = build_model()
strategy = tf.distribute.MirroredStrategy()
model.compile()

B. strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
model = build_model()

C. strategy = tf.distribute.MirroredStrategy()
model = build_model()
model.fit()

D. model = build_model()
model.fit()
strategy = tf.distribute.MirroredStrategy()

Solution

1. Step 1: Identify correct usage of strategy.scope()

   Model building and compiling must happen inside strategy.scope() to enable multi-GPU support.

2. Step 2: Check each option for correct order

   Only strategy = tf.distribute.MirroredStrategy() with strategy.scope(): model = build_model() builds the model inside the strategy scope. Others build or compile outside, which disables multi-GPU benefits.

3. Final Answer:

   strategy = tf.distribute.MirroredStrategy() with strategy.scope(): model = build_model() -> Option B

4. Quick Check:

   Build model inside strategy.scope() = strategy = tf.distribute.MirroredStrategy() with strategy.scope(): model = build_model() [OK]

Hint: Always build and compile inside strategy.scope() [OK]

Common Mistakes:

• Building model before strategy.scope()
• Compiling model outside strategy.scope()
• Calling fit before defining strategy

3. Given the code below, what will be the output shape of the predictions if the input batch size is 64 and there are 2 GPUs?

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

import numpy as np
inputs = np.random.random((64, 20))
predictions = model.predict(inputs)
print(predictions.shape)

medium

A. (10, 64)

B. (32, 10)

C. (128, 10)

D. (64, 10)

Solution

1. Step 1: Understand batch size and MirroredStrategy behavior

   MirroredStrategy splits the batch across GPUs internally but the output predictions shape matches the input batch size.

2. Step 2: Check model output shape

   The model outputs 10 classes per input. Input batch size is 64, so output shape is (64, 10).

3. Final Answer:

   (64, 10) -> Option D

4. Quick Check:

   Output shape = (batch_size, classes) = (64, 10) [OK]

Hint: Output shape matches input batch size, not multiplied by GPUs [OK]

Common Mistakes:

• Multiplying batch size by number of GPUs
• Confusing output shape dimensions
• Assuming output shape is (10, 64)

4. You wrote this code but get an error: RuntimeError: tf.distribute.MirroredStrategy() must be used with model creation inside its scope. What is the likely mistake?

strategy = tf.distribute.MirroredStrategy()
model = build_model()
with strategy.scope():
    model.compile(optimizer='adam', loss='mse')

medium

A. Model was compiled outside the strategy.scope()

B. MirroredStrategy cannot be used with compile()

C. Model was built outside the strategy.scope()

D. The optimizer 'adam' is not supported

Solution

1. Step 1: Analyze error message and code

   The error says model creation must be inside strategy.scope(). Here, model is built before entering the scope.

2. Step 2: Identify correct fix

   Building the model inside the scope fixes the error. Compiling inside scope is correct.

3. Final Answer:

   Model was built outside the strategy.scope() -> Option C

4. Quick Check:

   Build model inside scope to avoid RuntimeError [OK]

Hint: Build model inside strategy.scope() to fix this error [OK]

Common Mistakes:

• Building model before strategy.scope()
• Thinking compile() causes error
• Blaming optimizer choice

5. You want to train a large model on 4 GPUs using TensorFlow. Which of the following steps correctly applies multi-GPU training and ensures the model can handle larger batch sizes efficiently?

hard

A. Use tf.distribute.MirroredStrategy(), build and compile the model inside its scope, then increase batch size in model.fit()

B. Build the model first, then use tf.distribute.MirroredStrategy() only during model.fit() with default batch size

C. Use tf.distribute.MirroredStrategy() but keep batch size small to avoid memory issues

D. Train the model separately on each GPU without using MirroredStrategy and average weights manually

Solution

1. Step 1: Apply MirroredStrategy correctly

   Use tf.distribute.MirroredStrategy() and build/compile the model inside its scope to enable multi-GPU training.

2. Step 2: Adjust batch size for efficiency

   Increasing batch size in model.fit() leverages multiple GPUs to handle larger data efficiently.

3. Final Answer:

   Use tf.distribute.MirroredStrategy(), build and compile the model inside its scope, then increase batch size in model.fit() -> Option A

4. Quick Check:

   Build inside scope + increase batch size = efficient multi-GPU training [OK]

Hint: Build inside scope and increase batch size for multi-GPU speed [OK]

Common Mistakes:

• Building model outside strategy scope
• Not increasing batch size to use GPUs fully
• Manually averaging weights instead of using MirroredStrategy