如何在TensorFlow的Keras中配置模型并使用多TPU训练？

Hey there! I’ve messed around with multi-TPU training in Keras a bunch, so let me break down exactly how to set this up—since most examples default to single TPU addresses, it’s totally reasonable to be stuck here. Let’s go step by step.

Before diving into model code, you need to make sure TensorFlow can see all your TPUs. Use TPUClusterResolver to detect and connect to the cluster:

import tensorflow as tf

# For Cloud TPUs, this will auto-detect the cluster if you're in the same project/zone
# For on-prem or custom TPUs, pass comma-separated TPU addresses like:
# resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://tpu-node-1:8470,grpc://tpu-node-2:8470')
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()

# Initialize the TPU system—this sets up communication between TPUs
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
print(f"Number of TPUs available: {len(tf.config.list_logical_devices('TPU'))}")

If this prints a number greater than 1, you’re good to go!

Keras uses distribution strategies to handle multi-device training. For TPUs, we’ll use TPUStrategy:

strategy = tf.distribute.TPUStrategy(resolver)

This strategy will automatically split your model and data across all connected TPUs behind the scenes.

This is critical—every part of model creation (layers, optimizer, loss) needs to happen within the strategy’s scope. If you build the model outside, it won’t be distributed properly. Here’s an example:

with strategy.scope():
    # Define your model architecture
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(28,28,1)),
        tf.keras.layers.MaxPooling2D((2,2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])

    # Compile with a TPU-compatible optimizer/loss
    model.compile(
        optimizer=tf.keras.optimizers.Adam(),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
        metrics=['accuracy']
    )

Pro tip: Avoid using legacy optimizers (like tf.keras.optimizers.legacy.Adam) unless absolutely necessary—they’re less reliable with multi-TPU setups.

Multi-TPUs need data fed fast enough to keep all cores busy. Use tf.data.Dataset with these best practices:

• Scale your batch size: The global batch size should be per_tpu_core_batch_size * number_of_tpu_cores. For example, if each TPU core handles 64 samples, and you have 8 cores (common for a single TPU pod slice), your global batch size is 512.
• Use prefetching: Add .prefetch(tf.data.AUTOTUNE) to your dataset pipeline to overlap data loading and model execution.
• Avoid small datasets: If your dataset is tiny, the TPUs might sit idle between batches—consider augmenting data or using a larger dataset.

Example data pipeline:

def load_dataset():
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    x_train = x_train.reshape(-1, 28,28,1).astype('float32') / 255.0
    x_test = x_test.reshape(-1,28,28,1).astype('float32') /255.0

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    train_dataset = train_dataset.shuffle(10000).batch(512).prefetch(tf.data.AUTOTUNE)

    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
    test_dataset = test_dataset.batch(512).prefetch(tf.data.AUTOTUNE)

    return train_dataset, test_dataset

train_ds, test_ds = load_dataset()

Now just call model.fit() like you would with a single TPU or GPU—TPUStrategy handles all the distribution behind the scenes:

history = model.fit(
    train_ds,
    epochs=10,
    validation_data=test_ds
)

When saving the model, do it inside the strategy scope if you plan to load it back for multi-TPU use later:

with strategy.scope():
    model.save('multi_tpu_model.h5')  # Or use SavedModel format with model.save('model_dir')

To load it, wrap the load_model call in the strategy scope too.

• TPU initialization fails: Double-check your TPU addresses (if using custom clusters) and ensure network connectivity between your training machine and TPUs. For Cloud TPUs, confirm your VM is in the same zone as the TPU pod.
• OOM errors: Reduce the per-core batch size. Remember, the global batch size is multiplied by the number of cores, so even a small per-core size can add up.
• Model not using all TPUs: Make 100% sure your model is built and compiled inside the strategy.scope()—this is the most common mistake!
• Slow training: Check your data pipeline—if it’s not prefetching or batching correctly, the TPUs will wait for data. Use tf.data.AUTOTUNE to optimize prefetching.

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