TensorFlow 中如何使用 tf.distribute.MirroredStrategy？

推荐答案

// www.javascriptcn.com code example
import tensorflow as tf

# 创建 MirroredStrategy 对象
strategy = tf.distribute.MirroredStrategy()

# 在策略范围内定义模型
with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
        tf.keras.layers.Dense(10, activation='softmax')
    ])

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

# 加载数据
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 784).astype('float32') / 255
x_test = x_test.reshape(-1, 784).astype('float32') / 255

# 训练模型
model.fit(x_train, y_train, epochs=5, batch_size=64)

# 评估模型
model.evaluate(x_test, y_test)

本题详细解读

1. MirroredStrategy 简介

tf.distribute.MirroredStrategy 是 TensorFlow 提供的一种同步分布式训练策略，适用于单机多 GPU 环境。它会在每个 GPU 上复制模型的副本，并在每个步骤中同步更新模型参数。

2. 使用步骤

1. 创建 MirroredStrategy 对象：

   strategy = tf.distribute.MirroredStrategy()

   这会自动检测可用的 GPU 并创建相应的副本。

2. 在策略范围内定义模型：

   with strategy.scope():
       model = tf.keras.Sequential([
           tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
           tf.keras.layers.Dense(10, activation='softmax')
       ])

   在 strategy.scope() 范围内定义的模型会自动在每个 GPU 上复制。

3. 编译模型：

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

   编译模型时，优化器和损失函数会自动适应分布式训练。

4. 加载数据：

   (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
   x_train = x_train.reshape(-1, 784).astype('float32') / 255
   x_test = x_test.reshape(-1, 784).astype('float32') / 255

   数据加载和预处理与单 GPU 训练相同。

5. 训练模型：

   model.fit(x_train, y_train, epochs=5, batch_size=64)

   训练时，数据会自动分配到各个 GPU 上进行并行计算。

6. 评估模型：

   model.evaluate(x_test, y_test)

   评估时，模型会在所有 GPU 上同步计算结果。

3. 注意事项

• GPU 数量：MirroredStrategy 会自动检测可用的 GPU 数量，并创建相应数量的模型副本。
• 数据并行：每个 GPU 都会处理一部分数据，并在每个步骤结束时同步梯度。
• 性能优化：可以通过调整 batch_size 和 num_gpus 来优化训练性能。