How to set a TensorFlow distribution strategy

TensorFlow distribution strategies control where model variables are created and how each training step is replicated across visible devices. Setting the strategy before building a Keras model keeps the same training code usable on a CPU-only machine, one GPU, or a single host with multiple GPUs.

For one machine, tf.distribute.MirroredStrategy is the usual synchronous strategy. It creates a replica on each selected device, keeps variables mirrored, and combines updates after each batch; on a CPU-only host the same API still reports one synchronized replica.

Objects that create TensorFlow variables determine whether the strategy is actually used. Create the strategy first, open strategy.scope(), and instantiate the model, optimizer, metrics, and any checkpoint restore that creates variables inside that scope before starting model.fit() or a training-batch smoke test.

1. Save a minimal strategy check as

   strategy-demo.py

   .

   strategy-demo.py

   import tensorflow as tf
    
   strategy = tf.distribute.MirroredStrategy()
    
   with strategy.scope():
       inputs = tf.keras.Input(shape=(4,))
       outputs = tf.keras.layers.Dense(1)(inputs)
       model = tf.keras.Model(inputs, outputs)
       model.compile(optimizer="adam", loss="mse")
    
   features = tf.ones((8, 4))
   labels = tf.ones((8, 1))
   loss = model.train_on_batch(features, labels)
    
   print("strategy", strategy.__class__.__name__)
   print("replicas", strategy.num_replicas_in_sync)
   print("model_strategy", model.distribute_strategy is strategy)
   print("train_on_batch", loss >= 0)

   The model and optimizer are created while the strategy scope is active, and the training-batch call confirms the scoped model can execute one update.

2. Create the same strategy before building the real model in the project.

   strategy = tf.distribute.MirroredStrategy()
    
   with strategy.scope():
       model = build_model()
       model.compile(
           optimizer="adam",
           loss="mse",
           metrics=["mae"],
       )

   Restore checkpoints inside the same scope when the restore path creates variables on first use.
   Related: How to compile a model in Keras

3. Batch the training and validation datasets with the global batch size.

   global_batch = 256
    
   train_ds = make_train_dataset().batch(global_batch).prefetch(tf.data.AUTOTUNE)
   val_ds = make_validation_dataset().batch(global_batch).prefetch(tf.data.AUTOTUNE)

   The global batch is the total batch across all synchronized replicas. With four replicas and a global batch of 256, each replica receives 64 examples per step.

4. Call model.fit() after the model has captured the strategy.

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

   Keras stores the distribution context on the model, so the fit call can run outside the with strategy.scope() block. Custom training loops need additional strategy.run() handling.
   Related: How to run a custom training loop in TensorFlow

5. Restrict MirroredStrategy to specific devices only when the job should not use every visible accelerator.

   strategy = tf.distribute.MirroredStrategy(
       devices=[
           "/gpu:0",
           "/gpu:1",
       ],
   )

   Omitting devices= lets MirroredStrategy use every GPU that TensorFlow can already see; a CPU-only or single-GPU host still reports one synchronized replica.
   Related: How to enable GPU acceleration in TensorFlow

6. Run the strategy check and confirm the selected strategy is attached before the smoke-test training batch completes.

   $ python strategy-demo.py
   strategy MirroredStrategy
   replicas 1
   model_strategy True
   train_on_batch True

   A host with multiple visible GPUs prints a larger replicas value. Keep model_strategy True because it shows the model was created under the selected strategy.