keras开发者文档 11：多GPU和分布式训练

介绍

通常，有两种方法可以在多个设备之间分配计算：

数据并行性，其中一个模型可以在多个设备或多台机器上复制。它们每个处理不同批次的数据，然后合并结果。此设置存在许多变体，不同的模型副本合并结果的方式不同，它们是否在每个批次中保持同步，或者它们是否松散耦合等等。

模型并行性，其中单个模型的不同部分在不同的设备上运行，一起处理一批数据。这最适合具有自然并行架构的模型，例如具有多个分支的模型。

本指南重点介绍数据并行性，尤其是同步数据并行性，在该模型中，模型的不同副本在每次处理批次后均保持同步。同步性使模型收敛行为与单设备训练中看到的一致。

具体来说，本指南教您如何使用tf.distribute API在以下两种设置中以最少的代码更改就在多个GPU上训练Keras模型：

• 在单台计算机上安装的多个GPU（通常为2至8个）上（单主机，多设备训练）。对于研究人员和小型行业工作流程，这是最常见的设置。
• 在许多计算机的群集上，每台计算机都托管一个或多个GPU（多工作分布式训练）。这是大型行业工作流程的理想设置，例如使用20-100个GPU在数千万个图像上训练高分辨率图像分类模型。

设置

import tensorflow as tf
from tensorflow import keras

单主机，多设备同步训练

在此设置中，您有一台装有多个GPU（通常为2到8个）的计算机。每个设备都将运行模型的副本（称为副本）。为了简单起见，在下文中，我们假设我们正在处理8个GPU，而不会失去一般性。

它是怎么运作的

在训练的每个步骤：

• 当前的数据批次（称为global batch）分为8个不同的子批次（称为local batches）。例如，如果全局批次具有512个样本，则8个本地批次中的每个批次将具有64个样本。
• 8个副本中的每个副本均独立处理本地批处理：它们运行前向遍历，然后进行向后遍历，输出相对于本地批次模型损失的权重梯度。
• 源自局部梯度的权重更新将有效地合并到8个副本中。因为这是在每个步骤的最后完成的，所以副本始终保持同步。
  实际上，同步更新模型副本的权重的过程是在每个权重变量的级别上进行的。这是通过 mirrored variable对象完成的。

怎样去使用它

要使用Keras模型进行单主机，多设备同步训练，请使用tf.distribute.MirroredStrategy API 。 运作方式如下：

• 实例化MirroredStrategy，可以选择配置要使用的特定设备（默认情况下，该策略将使用所有可用的GPU）。
• 使用策略对象打开一个范围，并在此范围内创建您需要的所有包含变量的Keras对象。 通常，这意味着在分发范围内 creating & compiling the model。
• 像往常一样通过fit（）训练模型。
  重要的是，我们建议您使用tf.data.Dataset对象在多设备或分布式工作流中加载数据。

从示意图上看，它看起来像这样：

# Create a MirroredStrategy.
strategy = tf.distribute.MirroredStrategy()
print('Number of devices: {}'.format(strategy.num_replicas_in_sync))

# Open a strategy scope.
with strategy.scope():
  # Everything that creates variables should be under the strategy scope.
  # In general this is only model construction & `compile()`.
  model = Model(...)
  model.compile(...)

# Train the model on all available devices.
model.fit(train_dataset, validation_data=val_dataset, ...)

# Test the model on all available devices.
model.evaluate(test_dataset)

这是一个简单的端到端可运行示例：

def get_compiled_model():
    # Make a simple 2-layer densely-connected neural network.
    inputs = keras.Input(shape=(784,))
    x = keras.layers.Dense(256, activation="relu")(inputs)
    x = keras.layers.Dense(256, activation="relu")(x)
    outputs = keras.layers.Dense(10)(x)
    model = keras.Model(inputs, outputs)
    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=[keras.metrics.SparseCategoricalAccuracy()],
    )
    return model


def get_dataset():
    batch_size = 32
    num_val_samples = 10000

    # Return the MNIST dataset in the form of a `tf.data.Dataset`.
    (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

    # Preprocess the data (these are Numpy arrays)
    x_train = x_train.reshape(-1, 784).astype("float32") / 255
    x_test = x_test.reshape(-1, 784).astype("float32") / 255
    y_train = y_train.astype("float32")
    y_test = y_test.astype("float32")

    # Reserve num_val_samples samples for validation
    x_val = x_train[-num_val_samples:]
    y_val = y_train[-num_val_samples:]
    x_train = x_train[:-num_val_samples]
    y_train = y_train[:-num_val_samples]
    return (
        tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size),
        tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(batch_size),
        tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size),
    )


# Create a MirroredStrategy.
strategy = tf.distribute.MirroredStrategy()
print("Number of devices: {}".format(strategy.num_replicas_in_sync))

# Open a strategy scope.
with strategy.scope():
    # Everything that creates variables should be under the strategy scope.
    # In general this is only model construction & `compile()`.
    model = get_compiled_model()

# Train the model on all available devices.
train_dataset, val_dataset, test_dataset = get_dataset()
model.fit(train_dataset, epochs=2, validation_data=val_dataset)

# Test the model on all available devices.
model.evaluate(test_dataset)

使用回调确保容错

使用分布式训练时，应始终确保具有从故障中恢复的策略（容错）。 处理此问题的最简单方法是将ModelCheckpoint回调传递给fit（），以定期保存模型（例如，每100个批次或每个纪元）。 然后，您可以从保存的模型重新开始训练。

这是一个简单的例子：

import os
from tensorflow import keras

# Prepare a directory to store all the checkpoints.
checkpoint_dir = "./ckpt"
if not os.path.exists(checkpoint_dir):
    os.makedirs(checkpoint_dir)


def make_or_restore_model():
    # Either restore the latest model, or create a fresh one
    # if there is no checkpoint available.
    checkpoints = [checkpoint_dir + "/" + name for name in os.listdir(checkpoint_dir)]
    if checkpoints:
        latest_checkpoint = max(checkpoints, key=os.path.getctime)
        print("Restoring from", latest_checkpoint)
        return keras.models.load_model(latest_checkpoint)
    print("Creating a new model")
    return get_compiled_model()


def run_training(epochs=1):
    # Create a MirroredStrategy.
    strategy = tf.distribute.MirroredStrategy()

    # Open a strategy scope and create/restore the model
    with strategy.scope():
        model = make_or_restore_model()

    callbacks = [
        # This callback saves a SavedModel every epoch
        # We include the current epoch in the folder name.
        keras.callbacks.ModelCheckpoint(
            filepath=checkpoint_dir + "/ckpt-{epoch}", save_freq="epoch"
        )
    ]
    model.fit(
        train_dataset,
        epochs=epochs,
        callbacks=callbacks,
        validation_data=val_dataset,
        verbose=2,
    )


# Running the first time creates the model
run_training(epochs=1)

# Calling the same function again will resume from where we left off
run_training(epochs=1)

tf.data性能提示

在进行分布式训练时，加载数据的效率通常变得至关重要。这里有一些技巧，以确保您的tf.data管道尽可能快地运行。

关于数据集批处理的注意事项

创建数据集时，请确保使用全局批处理大小进行批处理。例如，如果您的8个GPU中的每一个都能够运行64个样本的批次，则您可以使用512的全局批次大小。

调用dataset.cache（）

如果对数据集调用.cache（），则在对数据进行第一次迭代后，将缓存其数据。每个后续迭代将使用缓存的数据。缓存可以在内存中（默认），也可以在您指定的本地文件中。

在以下情况下，可以提高性能：

• 您的数据不会因迭代而改变
• 您正在从远程分布式文件系统读取数据
• 您正在从本地磁盘读取数据，但是您的数据将容纳在内存中，并且您的工作流受到IO的限制（例如，读取和解码图像文件）。

调用dataset.prefetch（buffer_size）

创建数据集后，几乎应该总是调用.prefetch（buffer_size）。这意味着您的数据管道将与模型异步运行，新样本将被预处理并存储在缓冲区中，而当前批次样本将用于训练模型。当前批处理结束时，下一批将在GPU内存中预取。

多机分布式同步培训

这个怎么运作

在此设置中，您有多台计算机（称为工作器），每台计算机上都具有一个或多个GPU。就像单主机训练发生的情况一样，每个可用的GPU将运行一个模型副本，并且每个副本之后每个副本的变量值都保持同步。

重要的是，当前实现假定所有工作程序都具有相同数量的GPU（同质集群）。

如何使用它

1. 设置集群（我们在下面提供了指针）。
2. 在每个工作机上设置适当的TF_CONFIG环境变量。这告诉工作机其角色是什么以及如何与同事交流。
3. 在每位工作机上，在MultiWorkerMirroredStrategy object的范围内运行模型构建和编译代码，这与我们对单主机培训所做的类似。
4. 在指定的评估程序机器上运行评估代码。

设置集群

首先，建立一个集群（计算机的集合）。应该分别设置每台机器，以便能够运行您的模型（通常，每台机器将运行相同的Docker映像）并能够访问您的数据源（例如GCS）。

群集管理不在本指南的范围内。这是一个文档，可以帮助您入门。您也可以看看Kubeflow。

设置TF_CONFIG环境变量

尽管每个工作程序上运行的代码与单主机工作流程中使用的代码几乎相同（除了使用不同的tf.distribute策略对象），但单主机工作流程与多主机工作流程之间的一个重要区别是您需要在集群中运行的每台计算机上设置TF_CONFIG环境变量。

TF_CONFIG环境变量是一个JSON字符串，它指定：

• 集群配置，而组成集群的机器的地址和端口列表
• 工作机的“任务”，这是该特定计算机在集群中必须扮演的角色。
  TF_CONFIG的一个示例是：

os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        'worker': ["localhost:12345", "localhost:23456"]
    },
    'task': {'type': 'worker', 'index': 0}
})

在多机同步培训设置中，机器的有效角色（任务类型）是“工人”和“评估者”。

例如，如果您有8台计算机，每台计算机有4个GPU，则可以有7个工作人员和1个评估程序。

• 工人训练模型，每个模型处理一个全局批次的子批次。
• 其中一名工作人员（工作人员0）将充当“首席”，这是一种特殊的工作人员，负责保存日志和检查点以供以后重用（通常到云存储位置）。
• 评估程序运行一个连续循环，该循环加载首席工作人员保存的最新检查点，对其进行评估（与其他工作人员异步）并写入评估日志（例如TensorBoard日志）。

在每个工作程序上运行代码

您将在每个工人（包括主管）上运行训练代码，并在评估者上运行评估代码。

除了使用MultiWorkerMirroredStrategy而不是MirroredStrategy之外，训练代码与您在单主机设置中使用的代码基本相同。

每个工作程序将运行相同的代码（减去下面注释中说明的差异），包括相同的回调。

**注意：**保存模型检查点或日志的回调应为每个工作人员保存到不同的目录。按照惯例，所有工作人员都应保存到本地磁盘（通常是临时的），但工作人员0除外，这会将TensorBoard日志检查点保存到Cloud存储位置，以供以后访问和重用。

评估者将只使用MirroredStrategy（因为它在单台机器上运行并且不需要与其他机器进行通信）并调用model.evaluate（）。它将把首席工作人员保存的最新检查点加载到Cloud存储位置，并将评估日志保存到与首席日志相同的位置。

示例：在多机设置中运行的代码

在负责人（工人0）上：

# Set TF_CONFIG
os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        'worker': ["localhost:12345", "localhost:23456"]
    },
    'task': {'type': 'worker', 'index': 0}
})


# Open a strategy scope and create/restore the model.
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
with strategy.scope():
  model = make_or_restore_model()

callbacks = [
    # This callback saves a SavedModel every 100 batches
    keras.callbacks.ModelCheckpoint(filepath='path/to/cloud/location/ckpt',
                                    save_freq=100),
    keras.callbacks.TensorBoard('path/to/cloud/location/tb/')
]
model.fit(train_dataset,
          callbacks=callbacks,
          ...)

在其它工作者上：

# Set TF_CONFIG
worker_index = 1  # For instance
os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        'worker': ["localhost:12345", "localhost:23456"]
    },
    'task': {'type': 'worker', 'index': worker_index}
})


# Open a strategy scope and create/restore the model.
# You can restore from the checkpoint saved by the chief.
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
with strategy.scope():
  model = make_or_restore_model()

callbacks = [
    keras.callbacks.ModelCheckpoint(filepath='local/path/ckpt', save_freq=100),
    keras.callbacks.TensorBoard('local/path/tb/')
]
model.fit(train_dataset,
          callbacks=callbacks,
          ...)

在检测者上：

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
  model = make_or_restore_model()  # Restore from the checkpoint saved by the chief.

results = model.evaluate(val_dataset)
# Then, log the results on a shared location, write TensorBoard logs, etc

扩展阅读

1. TensorFlow distributed training guide.
2. Tutorial on multi-worker training with Keras.
3. MirroredStrategy docs.
4. MultiWorkerMirroredStrategy docs.

参考文献

Multi-GPU and distributed training