tf.distribute Strategies

Two lines of code, multi-GPU or TPU

TF's distribution API wraps a single-device script to run across multiple GPUs or TPU chips with minimal code changes. Create a strategy, build and compile inside strategy.scope(), and model.fit handles gradient synchronization automatically.

Strategy	Hardware	Sync	Status
MirroredStrategy	1 machine, N GPUs	Sync (NCCL)	Stable
TPUStrategy	TPU pod / v2/v3/v4	Sync	Stable
MultiWorkerMirroredStrategy	N machines × N GPUs	Sync	Stable
ParameterServerStrategy	N workers + param servers	Async	Experimental

Scale the global batch size with the number of replicas. A per-replica batch of 64 across 4 GPUs is a global batch of 256. Match learning rate accordingly (linear scaling rule, plus warmup for very large batches).

Code

MirroredStrategy — single machine, multi-GPU·python

import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()
print(f"Replicas: {strategy.num_replicas_in_sync}")

with strategy.scope():
    model = tf.keras.applications.ResNet50(weights=None, classes=1000)
    model.compile(
        optimizer=tf.keras.optimizers.Adam(1e-3),
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy'],
    )

PER_REPLICA = 64
GLOBAL_BATCH = PER_REPLICA * strategy.num_replicas_in_sync

model.fit(train_ds, epochs=10, validation_data=val_ds)

TPUStrategy — Google Cloud TPU 또는 Colab·python

import tensorflow as tf

# In Colab: empty TPU spec auto-detects
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)   # wipes TPU memory

strategy = tf.distribute.TPUStrategy(resolver)
print(f"TPU cores: {strategy.num_replicas_in_sync}")

with strategy.scope():
    model = tf.keras.applications.EfficientNetV2S(
        input_shape=(224, 224, 3), classes=10, weights=None,
    )
    model.compile(
        optimizer=tf.keras.optimizers.Adam(1e-3),
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy'],
    )

# TPU efficiency: large batches per core
BATCH_SIZE = 128 * strategy.num_replicas_in_sync
# All data must be tf.data.Dataset, not numpy arrays
model.fit(train_dataset, validation_data=val_dataset, epochs=20)

Two lines of code, multi-GPU or TPU

Code

Progress

Comments 0