使用 TensorFlow 进行分布式训练

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源码

下载笔记本

概览

tf.distribute.Strategy 是一个 TensorFlow API，用于在多个 GPU、多台机器或 TPU 上分发训练。使用此 API，您只需极少量的代码更改即可分发您现有的模型和训练代码。

tf.distribute.Strategy 的设计旨在实现以下关键目标：

易于使用，并支持多种用户群体，包括研究人员、机器学习工程师等。
开箱即用，提供良好的性能。
策略之间易于切换。

您可以使用 tf.distribute.Strategy 通过高级 API（如 Keras Model.fit）以及自定义训练循环（通常是任何使用 TensorFlow 的计算）来分发训练。

在 TensorFlow 2.x 中，您可以使用 tf.function 以动态图（eagerly）或图模式执行程序。tf.distribute.Strategy 旨在支持这两种执行模式，但与 tf.function 配合使用效果最佳。动态图模式仅建议用于调试目的，且不支持 tf.distribute.TPUStrategy。虽然本指南重点介绍训练，但此 API 也可用于在不同平台上分发评估和预测任务。

您可以非常轻松地更改代码以使用 tf.distribute.Strategy，因为 TensorFlow 的底层组件已更改为具有策略感知能力。这包括变量、层、模型、优化器、指标、摘要和检查点。

在本指南中，您将了解各种类型的策略以及如何在不同情况下使用它们。要了解如何调试性能问题，请查看优化 TensorFlow GPU 性能指南。

设置 TensorFlow

import tensorflow as tf

2024-10-25 03:10:09.809713: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1729825809.832772  192915 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1729825809.839425  192915 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

策略类型

tf.distribute.Strategy 旨在涵盖不同维度下的多种用例。其中一些组合目前已得到支持，其他组合将在未来加入。部分维度包括：

同步训练与异步训练： 这是使用数据并行进行分布式训练的两种常见方式。在同步训练中，所有工作节点在每一步都会同步对输入数据的不同切片进行训练，并聚合梯度。在异步训练中，所有工作节点独立地对输入数据进行训练并异步更新变量。通常，同步训练通过 all-reduce 实现，异步训练通过参数服务器架构实现。
硬件平台： 您可能希望将训练扩展到单机上的多个 GPU、网络中的多台机器（每台机器有 0 个或多个 GPU），或者云 TPU 上。

为了支持这些用例，TensorFlow 提供了 MirroredStrategy、TPUStrategy、MultiWorkerMirroredStrategy、ParameterServerStrategy、CentralStorageStrategy 以及其他可用策略。下一节将解释在 TensorFlow 中哪些场景支持这些策略。以下是快速概览：

训练 API	`MirroredStrategy`	`TPUStrategy`	`MultiWorkerMirroredStrategy`	`CentralStorageStrategy`	`ParameterServerStrategy`
Keras `Model.fit`	受支持	受支持	受支持	实验性支持	实验性支持
自定义训练循环	受支持	受支持	受支持	实验性支持	实验性支持
Estimator API	有限支持	不支持	有限支持	有限支持	有限支持

MirroredStrategy

tf.distribute.MirroredStrategy 支持在单机上的多个 GPU 上进行同步分布式训练。它为每个 GPU 设备创建一个副本。模型中的每个变量都会在所有副本上进行镜像。这些变量共同构成一个单一的逻辑变量，称为 MirroredVariable。这些变量通过应用相同的更新来保持同步。

高效的 all-reduce 算法被用于在设备间通信变量更新。All-reduce 通过相加聚合所有设备上的张量，并使它们在每个设备上都可用。这是一种融合算法，非常高效，可以显著降低同步开销。根据设备间可用的通信类型，有许多 all-reduce 算法和实现可用。默认情况下，它使用 NVIDIA 集体通信库 (NCCL) 作为 all-reduce 的实现。您可以从其他几个选项中进行选择，或编写自己的实现。

这是创建 MirroredStrategy 的最简单方法：

mirrored_strategy = tf.distribute.MirroredStrategy()

INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)
W0000 00:00:1729825812.490898  192915 gpu_device.cc:2344] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://tensorflowcn.cn/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...

这将创建一个 MirroredStrategy 实例，它将使用 TensorFlow 可见的所有 GPU，并使用 NCCL 作为跨设备通信。

如果您只想使用机器上的部分 GPU，可以这样做：

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

INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1')

如果您想覆盖跨设备通信，可以使用 cross_device_ops 参数并提供一个 tf.distribute.CrossDeviceOps 实例。目前，除了默认的 tf.distribute.NcclAllReduce 之外，tf.distribute.HierarchicalCopyAllReduce 和 tf.distribute.ReductionToOneDevice 是另外两个可选方案。

mirrored_strategy = tf.distribute.MirroredStrategy(
    cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())

INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)

TPUStrategy

tf.distribute.TPUStrategy 让您可以在张量处理单元 (TPUs) 上运行 TensorFlow 训练。TPU 是 Google 的专用 ASIC，旨在显著加速机器学习工作负载。它们可在 Google Colab、TPU Research Cloud 和 Cloud TPU 上使用。

在分布式训练架构方面，TPUStrategy 与 MirroredStrategy 相同——它实现了同步分布式训练。TPU 提供了自己高效的 all-reduce 以及跨多个 TPU 核心的其他集合操作实现，这些实现被用于 TPUStrategy 中。

这是实例化 TPUStrategy 的方法：

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu=tpu_address)
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
tpu_strategy = tf.distribute.TPUStrategy(cluster_resolver)

TPUClusterResolver 实例有助于定位 TPU。在 Colab 中，您无需为其指定任何参数。

如果您想将其用于 Cloud TPU：

您必须在 tpu 参数中指定 TPU 资源的名称。
您必须在程序开始时显式初始化 TPU 系统。这是在使用 TPU 进行计算之前所必需的。初始化 TPU 系统也会擦除 TPU 内存，因此首先完成此步骤以避免丢失状态非常重要。

MultiWorkerMirroredStrategy

tf.distribute.MultiWorkerMirroredStrategy 与 MirroredStrategy 非常相似。它实现了跨多个工作节点的同步分布式训练，每个节点可能拥有多个 GPU。与 tf.distribute.MirroredStrategy 类似，它会在所有工作节点上的每个设备上创建模型中所有变量的副本。

这是创建 MultiWorkerMirroredStrategy 的最简单方法：

strategy = tf.distribute.MultiWorkerMirroredStrategy()

WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/device:CPU:0',)
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:CPU:0',), communication = CommunicationImplementation.AUTO

MultiWorkerMirroredStrategy 具有两种跨设备通信实现。CommunicationImplementation.RING 基于 RPC 并支持 CPU 和 GPU。CommunicationImplementation.NCCL 使用 NCCL 并在 GPU 上提供一流的性能，但不支持 CPU。CollectiveCommunication.AUTO 将选择权留给 TensorFlow。您可以按以下方式指定它们：

communication_options = tf.distribute.experimental.CommunicationOptions(
    implementation=tf.distribute.experimental.CommunicationImplementation.NCCL)
strategy = tf.distribute.MultiWorkerMirroredStrategy(
    communication_options=communication_options)

WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/device:CPU:0',)
WARNING:tensorflow:Enabled NCCL communication but no GPUs detected/specified.
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:CPU:0',), communication = CommunicationImplementation.NCCL

与多 GPU 训练相比，实现多工作节点训练的一个关键区别是多工作节点设置。'TF_CONFIG' 环境变量是 TensorFlow 中为集群中每个工作节点指定集群配置的标准方式。请在本文档的设置 TF_CONFIG 章节中了解更多信息。

有关 MultiWorkerMirroredStrategy 的更多详细信息，请参考以下教程：

ParameterServerStrategy

参数服务器训练是一种常用的数据并行方法，用于扩展多台机器上的模型训练。参数服务器训练集群由工作节点和参数服务器组成。变量在参数服务器上创建，并在每一步由工作节点读取和更新。查看参数服务器训练教程以获取详细信息。

在 TensorFlow 2 中，参数服务器训练通过 tf.distribute.experimental.coordinator.ClusterCoordinator 类使用基于中心协调器的架构。

在此实现中，worker（工作节点）和 parameter server（参数服务器）任务运行 tf.distribute.Server，它们监听来自协调器的任务。协调器负责创建资源、调度训练任务、写入检查点并处理任务故障。

在协调器上运行的程序中，您将使用 ParameterServerStrategy 对象来定义训练步长，并使用 ClusterCoordinator 将训练步长调度到远程工作节点。这是创建它们的最简单方法：

strategy = tf.distribute.experimental.ParameterServerStrategy(
    tf.distribute.cluster_resolver.TFConfigClusterResolver(),
    variable_partitioner=variable_partitioner)
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
    strategy)

要了解有关 ParameterServerStrategy 的更多信息，请查看使用 Keras Model.fit 和自定义训练循环进行参数服务器训练教程。

在 TensorFlow 1 中，ParameterServerStrategy 仅通过 tf.compat.v1.distribute.experimental.ParameterServerStrategy 符号并通过 Estimator 提供。

CentralStorageStrategy

tf.distribute.experimental.CentralStorageStrategy 也执行同步训练。变量不会被镜像，而是放置在 CPU 上，运算被复制到所有本地 GPU 上。如果只有一个 GPU，则所有变量和运算都将放置在该 GPU 上。

创建 CentralStorageStrategy 实例的方法：

central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()

INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ['/job:localhost/replica:0/task:0/device:CPU:0'], variable_device = '/job:localhost/replica:0/task:0/device:CPU:0'

这将创建一个 CentralStorageStrategy 实例，该实例将使用所有可见的 GPU 和 CPU。副本上的变量更新在应用于变量之前将被聚合。

其他策略

除上述策略外，还有两种策略在配合 tf.distribute API 进行原型设计和调试时非常有用。

默认策略

默认策略（Default Strategy）是一种在没有显式分发策略生效时存在的分配策略。它实现了 tf.distribute.Strategy 接口，但它是直通的，不提供实际的分发。例如，Strategy.run(fn) 将直接调用 fn。使用此策略编写的代码的表现应与不使用任何策略编写的代码完全一致。您可以将其视为一个“空操作（no-op）”策略。

默认策略是一个单例，不能创建更多实例。它可以在任何显式策略作用域之外通过 tf.distribute.get_strategy 获取（这与在显式策略作用域内获取当前策略的 API 相同）。

default_strategy = tf.distribute.get_strategy()

此策略主要有两个用途：

它允许无条件地编写感知分发的库代码。例如，在 tf.keras.optimizers 中，您可以使用 tf.distribute.get_strategy 并使用该策略来缩减梯度——它总是会返回一个可以在其上调用 Strategy.reduce API 的策略对象。

# In optimizer or other library code
# Get currently active strategy
strategy = tf.distribute.get_strategy()
strategy.reduce("SUM", 1., axis=None)  # reduce some values

1.0

与库代码类似，它可用于编写使最终用户的程序在有无分发策略的情况下均能运行，而无需条件逻辑。以下是一个展示此功能的示例代码片段：

if tf.config.list_physical_devices('GPU'):
  strategy = tf.distribute.MirroredStrategy()
else:  # Use the Default Strategy
  strategy = tf.distribute.get_strategy()

with strategy.scope():
  # Do something interesting
  print(tf.Variable(1.))

<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>

OneDeviceStrategy

tf.distribute.OneDeviceStrategy 是一种将所有变量和计算放置在单个指定设备上的策略。

strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")

此策略在多个方面与默认策略不同。在默认策略中，与在没有任何分发策略的情况下运行 TensorFlow 相比，变量放置逻辑保持不变。但在使用 OneDeviceStrategy 时，在其作用域内创建的所有变量都会被显式放置在指定的设备上。此外，通过 OneDeviceStrategy.run 调用的任何函数也将被放置在指定的设备上。

通过此策略分发的输入将被预取到指定的设备。在默认策略中，没有输入分发。

与默认策略类似，此策略也可以用于在切换到真正分发到多个设备/机器的其他策略之前测试您的代码。这会比默认策略更多地调用分发策略机制，但不会达到使用（例如）MirroredStrategy 或 TPUStrategy 的全部程度。如果您想要代码表现得就像没有策略一样，请使用默认策略。

到目前为止，您已经了解了不同的策略以及如何实例化它们。接下来的几节将展示您可以使用它们分发训练的不同方式。

将 tf.distribute.Strategy 与 Keras Model.fit 结合使用

tf.distribute.Strategy 已集成到 tf.keras 中，这是 TensorFlow 对 Keras API 规范的实现。tf.keras 是用于构建和训练模型的高级 API。通过集成到 tf.keras 后端，使用 Keras 训练框架编写的训练代码可以通过 Model.fit 无缝地进行分发。

以下是您在代码中需要进行的更改：

创建适当的 tf.distribute.Strategy 实例。
将 Keras 模型、优化器和指标的创建移动到 strategy.scope 中。这样，模型 call()、train_step() 和 test_step() 方法中的代码都将被分发并在加速器上执行。

TensorFlow 分发策略支持所有类型的 Keras 模型——Sequential、函数式 (Functional) 和子类化 (Subclassed) 模型。

以下是针对一个具有一个 Dense 层的极其简单的 Keras 模型进行此操作的代码片段：

mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  model = tf.keras.Sequential([
      tf.keras.layers.Dense(1, input_shape=(1,),
                            kernel_regularizer=tf.keras.regularizers.L2(1e-4))])
  model.compile(loss='mse', optimizer='sgd')

INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)
/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/layers/core/dense.py:87: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead.
  super().__init__(activity_regularizer=activity_regularizer, **kwargs)

此示例使用 MirroredStrategy，因此您可以在具有多个 GPU 的机器上运行它。strategy.scope() 向 Keras 指明使用哪种策略来分发训练。在此作用域内创建模型/优化器/指标允许您创建分布式变量而不是常规变量。设置完成后，您可以像平时一样拟合模型。MirroredStrategy 会负责在可用 GPU 上复制模型训练、聚合梯度等工作。

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)
model.fit(dataset, epochs=2)
model.evaluate(dataset)

2024-10-25 03:10:12.686928: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
Epoch 1/2
 1/10 ━━━━━━━━━━━━━━━━━━━━ 2s 287ms/step - loss: 0.649210/10 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - loss: 0.5412
Epoch 2/2
 1/10 ━━━━━━━━━━━━━━━━━━━━ 0s 61ms/step - loss: 0.286910/10 ━━━━━━━━━━━━━━━━━━━━ 0s 1ms/step - loss: 0.2392
 1/10 ━━━━━━━━━━━━━━━━━━━━ 1s 168ms/step - loss: 0.126810/10 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - loss: 0.1268
0.1268472820520401

这里 tf.data.Dataset 提供了训练和评估输入。您也可以使用 NumPy 数组：

import numpy as np

inputs, targets = np.ones((100, 1)), np.ones((100, 1))
model.fit(inputs, targets, epochs=2, batch_size=10)

Epoch 1/2
 2/10 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - loss: 0.124410/10 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - loss: 0.1058
Epoch 2/2
 3/10 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - loss: 0.0539 10/10 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - loss: 0.0468
<keras.src.callbacks.history.History at 0x7fab904bcfa0>

在两种情况下（使用 Dataset 或 NumPy），给定输入的每个批次都会被平均分配到多个副本中。例如，如果您使用带有 2 个 GPU 的 MirroredStrategy，每个大小为 10 的批次将分配给 2 个 GPU，每个 GPU 在每一步接收 5 个输入样本。随着添加更多 GPU，每个周期（epoch）的训练速度都会加快。通常，随着添加更多加速器，您需要增加批次大小，以便有效利用额外的计算能力。您还需要根据模型重新调整学习率。您可以使用 strategy.num_replicas_in_sync 获取副本数量。

mirrored_strategy.num_replicas_in_sync

# Compute a global batch size using a number of replicas.
BATCH_SIZE_PER_REPLICA = 5
global_batch_size = (BATCH_SIZE_PER_REPLICA *
                     mirrored_strategy.num_replicas_in_sync)
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100)
dataset = dataset.batch(global_batch_size)

LEARNING_RATES_BY_BATCH_SIZE = {5: 0.1, 10: 0.15, 20:0.175}
learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]

目前支持哪些功能？

训练 API	`MirroredStrategy`	`TPUStrategy`	`MultiWorkerMirroredStrategy`	`ParameterServerStrategy`	`CentralStorageStrategy`
Keras `Model.fit`	受支持	受支持	受支持	实验性支持	实验性支持

示例和教程

以下是一系列教程和示例，它们说明了上述与 Keras Model.fit 的集成全流程：

教程：使用 Model.fit 和 MirroredStrategy 进行训练。
教程：使用 Model.fit 和 MultiWorkerMirroredStrategy 进行训练。
指南：包含使用 Model.fit 和 TPUStrategy 的示例。
教程：使用 Model.fit 和 ParameterServerStrategy 进行参数服务器训练。
教程：使用 Model.fit 和 TPUStrategy 为 GLUE 基准测试的许多任务微调 BERT。
TensorFlow Model Garden 存储库，其中包含使用各种策略实现的最新模型集合。

将 tf.distribute.Strategy 与自定义训练循环结合使用

如上所示，将 tf.distribute.Strategy 与 Keras Model.fit 结合使用仅需更改几行代码。通过稍微多一点努力，您还可以使用 tf.distribute.Strategy 编写自定义训练循环。

如果您需要比 Estimator 或 Keras 允许的范围更大的训练循环灵活性和控制力，您可以编写自定义训练循环。例如，使用 GAN 时，您可能希望在每一轮中采用不同数量的生成器或判别器步长。同样，高级框架不太适合强化学习训练。

tf.distribute.Strategy 类提供了一组核心方法来支持自定义训练循环。使用它们可能最初需要对代码进行少许重构，但一旦完成，您只需更改策略实例，即可在 GPU、TPU 和多台机器之间轻松切换。

以下是一个简短的代码片段，说明了使用与之前相同的 Keras 模型进行简单训练的用例。

首先，在策略的作用域内创建模型和优化器。这确保了随模型和优化器创建的任何变量都是镜像变量。

with mirrored_strategy.scope():
  model = tf.keras.Sequential([
      tf.keras.layers.Dense(1, input_shape=(1,),
                            kernel_regularizer=tf.keras.regularizers.L2(1e-4))])
  optimizer = tf.keras.optimizers.SGD()

接下来，创建输入数据集并调用 tf.distribute.Strategy.experimental_distribute_dataset 以根据策略分发数据集。

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(1000).batch(
    global_batch_size)
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)

然后，定义训练的一步。使用 tf.GradientTape 计算梯度，并使用优化器应用这些梯度来更新模型的变量。要分发此训练步，请将其放入函数 train_step 中，并将其与之前创建的 dist_dataset 中的数据集输入一起传递给 tf.distribute.Strategy.run。

# Sets `reduction=NONE` to leave it to tf.nn.compute_average_loss() below.
loss_object = tf.keras.losses.BinaryCrossentropy(
  from_logits=True,
  reduction=tf.keras.losses.Reduction.NONE)

def train_step(inputs):
  features, labels = inputs

  with tf.GradientTape() as tape:
    predictions = model(features, training=True)
    per_example_loss = loss_object(labels, predictions)
    loss = tf.nn.compute_average_loss(per_example_loss)
    model_losses = model.losses
    if model_losses:
      loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))

  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return loss

@tf.function
def distributed_train_step(dist_inputs):
  per_replica_losses = mirrored_strategy.run(train_step, args=(dist_inputs,))
  return mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
                         axis=None)

上述代码中还有几点需要注意：

您使用了 tf.nn.compute_average_loss 将每个样本的预测损失缩减为标量。tf.nn.compute_average_loss 对每个样本的损失求和，并将总和除以全局批次大小。这一点很重要，因为在每个副本上计算梯度后，它们会通过求和在副本间进行聚合。

默认情况下，全局批次大小被视为 tf.get_strategy().num_replicas_in_sync * tf.shape(per_example_loss)[0]。它也可以作为关键字参数 global_batch_size= 显式指定。如果没有较短的批次（short batches），默认值等效于 tf.nn.compute_average_loss(..., global_batch_size=global_batch_size)（其中 global_batch_size 定义如上）。（有关较短批次以及如何避免或处理它们的更多信息，请参阅自定义训练教程。）
您使用了 tf.nn.scale_regularization_loss，以也将使用 Model 对象注册的正则化损失（如果有）按 1/num_replicas_in_sync 进行缩放。对于那些与输入相关的正则化损失，需要建模代码（而非训练循环）在每个副本的批次大小上执行平均；这样建模代码可以保持对复制的不可知，而训练循环则对正则化损失的计算方式保持不可知。
当您在分发策略作用域内调用 apply_gradients 时，其行为会发生改变。具体而言，在同步训练期间，在每个并行实例上应用梯度之前，它会执行梯度的跨副本求和。
您还使用了 tf.distribute.Strategy.reduce API 来聚合 tf.distribute.Strategy.run 返回的结果以进行报告。tf.distribute.Strategy.run 返回策略中每个本地副本的结果，有多种方法可以消耗此结果。您可以 reduce 它们以获得聚合值。您还可以执行 tf.distribute.Strategy.experimental_local_results 以获取结果中包含的值列表，每个本地副本一个值。

最后，定义训练步长后，您可以遍历 dist_dataset 并在循环中运行训练：

for dist_inputs in dist_dataset:
  print(distributed_train_step(dist_inputs))

tf.Tensor(0.9024367, shape=(), dtype=float32)
tf.Tensor(0.8953863, shape=(), dtype=float32)
tf.Tensor(0.8884038, shape=(), dtype=float32)
tf.Tensor(0.88148874, shape=(), dtype=float32)
tf.Tensor(0.87464076, shape=(), dtype=float32)
tf.Tensor(0.86785895, shape=(), dtype=float32)
tf.Tensor(0.86114323, shape=(), dtype=float32)
tf.Tensor(0.8544927, shape=(), dtype=float32)
tf.Tensor(0.84790725, shape=(), dtype=float32)
tf.Tensor(0.841386, shape=(), dtype=float32)
tf.Tensor(0.83492863, shape=(), dtype=float32)
tf.Tensor(0.8285344, shape=(), dtype=float32)
tf.Tensor(0.82220304, shape=(), dtype=float32)
tf.Tensor(0.8159339, shape=(), dtype=float32)
tf.Tensor(0.8097264, shape=(), dtype=float32)
tf.Tensor(0.8035801, shape=(), dtype=float32)
tf.Tensor(0.79749453, shape=(), dtype=float32)
tf.Tensor(0.79146886, shape=(), dtype=float32)
tf.Tensor(0.785503, shape=(), dtype=float32)
tf.Tensor(0.779596, shape=(), dtype=float32)
tf.Tensor(0.77374756, shape=(), dtype=float32)
tf.Tensor(0.7679571, shape=(), dtype=float32)
tf.Tensor(0.7622242, shape=(), dtype=float32)
tf.Tensor(0.7565481, shape=(), dtype=float32)
tf.Tensor(0.75092846, shape=(), dtype=float32)
tf.Tensor(0.7453647, shape=(), dtype=float32)
tf.Tensor(0.73985624, shape=(), dtype=float32)
tf.Tensor(0.7344028, shape=(), dtype=float32)
tf.Tensor(0.7290035, shape=(), dtype=float32)
tf.Tensor(0.723658, shape=(), dtype=float32)
tf.Tensor(0.7183659, shape=(), dtype=float32)
tf.Tensor(0.71312654, shape=(), dtype=float32)
tf.Tensor(0.7079393, shape=(), dtype=float32)
tf.Tensor(0.70280397, shape=(), dtype=float32)
tf.Tensor(0.6977197, shape=(), dtype=float32)
tf.Tensor(0.69268626, shape=(), dtype=float32)
tf.Tensor(0.687703, shape=(), dtype=float32)
tf.Tensor(0.68276954, shape=(), dtype=float32)
tf.Tensor(0.67788523, shape=(), dtype=float32)
tf.Tensor(0.6730496, shape=(), dtype=float32)
tf.Tensor(0.66826224, shape=(), dtype=float32)
tf.Tensor(0.66352266, shape=(), dtype=float32)
tf.Tensor(0.6588302, shape=(), dtype=float32)
tf.Tensor(0.6541846, shape=(), dtype=float32)
tf.Tensor(0.6495853, shape=(), dtype=float32)
tf.Tensor(0.64503175, shape=(), dtype=float32)
tf.Tensor(0.6405235, shape=(), dtype=float32)
tf.Tensor(0.6360602, shape=(), dtype=float32)
tf.Tensor(0.6316412, shape=(), dtype=float32)
tf.Tensor(0.62726617, shape=(), dtype=float32)
tf.Tensor(0.6229345, shape=(), dtype=float32)
tf.Tensor(0.61864597, shape=(), dtype=float32)
tf.Tensor(0.6143999, shape=(), dtype=float32)
tf.Tensor(0.6101959, shape=(), dtype=float32)
tf.Tensor(0.60603356, shape=(), dtype=float32)
tf.Tensor(0.60191244, shape=(), dtype=float32)
tf.Tensor(0.597832, shape=(), dtype=float32)
tf.Tensor(0.5937919, shape=(), dtype=float32)
tf.Tensor(0.5897917, shape=(), dtype=float32)
tf.Tensor(0.585831, shape=(), dtype=float32)
tf.Tensor(0.58190924, shape=(), dtype=float32)
tf.Tensor(0.5780261, shape=(), dtype=float32)
tf.Tensor(0.57418114, shape=(), dtype=float32)
tf.Tensor(0.57037395, shape=(), dtype=float32)
tf.Tensor(0.5666041, shape=(), dtype=float32)
tf.Tensor(0.56287116, shape=(), dtype=float32)
tf.Tensor(0.55917484, shape=(), dtype=float32)
tf.Tensor(0.5555145, shape=(), dtype=float32)
tf.Tensor(0.55189, shape=(), dtype=float32)
tf.Tensor(0.54830086, shape=(), dtype=float32)
tf.Tensor(0.54474664, shape=(), dtype=float32)
tf.Tensor(0.54122704, shape=(), dtype=float32)
tf.Tensor(0.5377416, shape=(), dtype=float32)
tf.Tensor(0.5342899, shape=(), dtype=float32)
tf.Tensor(0.5308717, shape=(), dtype=float32)
tf.Tensor(0.5274865, shape=(), dtype=float32)
tf.Tensor(0.52413404, shape=(), dtype=float32)
tf.Tensor(0.52081394, shape=(), dtype=float32)
tf.Tensor(0.51752573, shape=(), dtype=float32)
tf.Tensor(0.5142692, shape=(), dtype=float32)
tf.Tensor(0.51104385, shape=(), dtype=float32)
tf.Tensor(0.50784945, shape=(), dtype=float32)
tf.Tensor(0.50468564, shape=(), dtype=float32)
tf.Tensor(0.50155205, shape=(), dtype=float32)
tf.Tensor(0.49844825, shape=(), dtype=float32)
tf.Tensor(0.4953741, shape=(), dtype=float32)
tf.Tensor(0.49232918, shape=(), dtype=float32)
tf.Tensor(0.4893132, shape=(), dtype=float32)
tf.Tensor(0.48632562, shape=(), dtype=float32)
tf.Tensor(0.4833664, shape=(), dtype=float32)
tf.Tensor(0.4804351, shape=(), dtype=float32)
tf.Tensor(0.47753143, shape=(), dtype=float32)
tf.Tensor(0.47465506, shape=(), dtype=float32)
tf.Tensor(0.47180572, shape=(), dtype=float32)
tf.Tensor(0.46898302, shape=(), dtype=float32)
tf.Tensor(0.4661867, shape=(), dtype=float32)
tf.Tensor(0.46341658, shape=(), dtype=float32)
tf.Tensor(0.4606722, shape=(), dtype=float32)
tf.Tensor(0.4579534, shape=(), dtype=float32)
tf.Tensor(0.4552598, shape=(), dtype=float32)
tf.Tensor(0.45259115, shape=(), dtype=float32)
tf.Tensor(0.44994718, shape=(), dtype=float32)
tf.Tensor(0.44732755, shape=(), dtype=float32)
tf.Tensor(0.44473216, shape=(), dtype=float32)
tf.Tensor(0.44216052, shape=(), dtype=float32)
tf.Tensor(0.4396125, shape=(), dtype=float32)
tf.Tensor(0.43708783, shape=(), dtype=float32)
tf.Tensor(0.4345862, shape=(), dtype=float32)
tf.Tensor(0.4321074, shape=(), dtype=float32)
tf.Tensor(0.42965108, shape=(), dtype=float32)
tf.Tensor(0.4272171, shape=(), dtype=float32)
tf.Tensor(0.42480516, shape=(), dtype=float32)
tf.Tensor(0.42241505, shape=(), dtype=float32)
tf.Tensor(0.42004645, shape=(), dtype=float32)
tf.Tensor(0.41769922, shape=(), dtype=float32)
tf.Tensor(0.41537297, shape=(), dtype=float32)
tf.Tensor(0.41306767, shape=(), dtype=float32)
tf.Tensor(0.41078293, shape=(), dtype=float32)
tf.Tensor(0.4085186, shape=(), dtype=float32)
tf.Tensor(0.4062744, shape=(), dtype=float32)
tf.Tensor(0.4040502, shape=(), dtype=float32)
tf.Tensor(0.40184572, shape=(), dtype=float32)
tf.Tensor(0.39966068, shape=(), dtype=float32)
tf.Tensor(0.3974949, shape=(), dtype=float32)
tf.Tensor(0.39534825, shape=(), dtype=float32)
tf.Tensor(0.39322042, shape=(), dtype=float32)
tf.Tensor(0.39111122, shape=(), dtype=float32)
tf.Tensor(0.3890205, shape=(), dtype=float32)
tf.Tensor(0.38694802, shape=(), dtype=float32)
tf.Tensor(0.38489357, shape=(), dtype=float32)
tf.Tensor(0.38285697, shape=(), dtype=float32)
tf.Tensor(0.38083804, shape=(), dtype=float32)
tf.Tensor(0.3788365, shape=(), dtype=float32)
tf.Tensor(0.37685227, shape=(), dtype=float32)
tf.Tensor(0.3748851, shape=(), dtype=float32)
tf.Tensor(0.37293482, shape=(), dtype=float32)
tf.Tensor(0.37100127, shape=(), dtype=float32)
tf.Tensor(0.36908418, shape=(), dtype=float32)
tf.Tensor(0.36718345, shape=(), dtype=float32)
tf.Tensor(0.3652989, shape=(), dtype=float32)
tf.Tensor(0.36343032, shape=(), dtype=float32)
tf.Tensor(0.36157757, shape=(), dtype=float32)
tf.Tensor(0.35974047, shape=(), dtype=float32)
tf.Tensor(0.3579188, shape=(), dtype=float32)
tf.Tensor(0.35611248, shape=(), dtype=float32)
tf.Tensor(0.3543213, shape=(), dtype=float32)
tf.Tensor(0.35254508, shape=(), dtype=float32)
tf.Tensor(0.3507837, shape=(), dtype=float32)
tf.Tensor(0.34903696, shape=(), dtype=float32)
tf.Tensor(0.34730473, shape=(), dtype=float32)
tf.Tensor(0.3455869, shape=(), dtype=float32)
tf.Tensor(0.3438832, shape=(), dtype=float32)
tf.Tensor(0.34219357, shape=(), dtype=float32)
tf.Tensor(0.3405178, shape=(), dtype=float32)
tf.Tensor(0.3388558, shape=(), dtype=float32)
tf.Tensor(0.3372074, shape=(), dtype=float32)
tf.Tensor(0.33557245, shape=(), dtype=float32)
tf.Tensor(0.33395082, shape=(), dtype=float32)
tf.Tensor(0.33234236, shape=(), dtype=float32)
tf.Tensor(0.33074695, shape=(), dtype=float32)
tf.Tensor(0.32916442, shape=(), dtype=float32)
tf.Tensor(0.3275946, shape=(), dtype=float32)
tf.Tensor(0.3260375, shape=(), dtype=float32)
tf.Tensor(0.3244928, shape=(), dtype=float32)
tf.Tensor(0.3229605, shape=(), dtype=float32)
tf.Tensor(0.32144046, shape=(), dtype=float32)
tf.Tensor(0.31993246, shape=(), dtype=float32)
tf.Tensor(0.3184365, shape=(), dtype=float32)
tf.Tensor(0.31695238, shape=(), dtype=float32)
tf.Tensor(0.31548, shape=(), dtype=float32)
tf.Tensor(0.31401917, shape=(), dtype=float32)
tf.Tensor(0.3125699, shape=(), dtype=float32)
tf.Tensor(0.31113195, shape=(), dtype=float32)
tf.Tensor(0.30970532, shape=(), dtype=float32)
tf.Tensor(0.3082898, shape=(), dtype=float32)
tf.Tensor(0.30688527, shape=(), dtype=float32)
tf.Tensor(0.3054917, shape=(), dtype=float32)
tf.Tensor(0.30410892, shape=(), dtype=float32)
tf.Tensor(0.3027368, shape=(), dtype=float32)
tf.Tensor(0.30137527, shape=(), dtype=float32)
tf.Tensor(0.3000242, shape=(), dtype=float32)
tf.Tensor(0.29868355, shape=(), dtype=float32)
tf.Tensor(0.29735315, shape=(), dtype=float32)
tf.Tensor(0.29603288, shape=(), dtype=float32)
tf.Tensor(0.29472268, shape=(), dtype=float32)
tf.Tensor(0.2934224, shape=(), dtype=float32)
tf.Tensor(0.29213202, shape=(), dtype=float32)
tf.Tensor(0.29085135, shape=(), dtype=float32)
tf.Tensor(0.28958035, shape=(), dtype=float32)
tf.Tensor(0.2883189, shape=(), dtype=float32)
tf.Tensor(0.28706694, shape=(), dtype=float32)
tf.Tensor(0.28582436, shape=(), dtype=float32)
tf.Tensor(0.28459102, shape=(), dtype=float32)
tf.Tensor(0.28336692, shape=(), dtype=float32)
tf.Tensor(0.2821518, shape=(), dtype=float32)
tf.Tensor(0.28094578, shape=(), dtype=float32)
tf.Tensor(0.27974862, shape=(), dtype=float32)
tf.Tensor(0.2785603, shape=(), dtype=float32)
tf.Tensor(0.27738073, shape=(), dtype=float32)
tf.Tensor(0.2762098, shape=(), dtype=float32)

在上面的示例中，您遍历了 dist_dataset 以提供输入给您的训练。您还获得了 tf.distribute.Strategy.make_experimental_numpy_dataset 以支持 NumPy 输入。您可以在调用 tf.distribute.Strategy.experimental_distribute_dataset 之前使用此 API 创建数据集。

遍历数据的另一种方法是显式使用迭代器。当您想要运行特定步数而不是遍历整个数据集时，您可能希望这样做。上述迭代现在将被修改为首先创建一个迭代器，然后显式调用其上的 next 以获取输入数据。

iterator = iter(dist_dataset)
for _ in range(10):
  print(distributed_train_step(next(iterator)))

tf.Tensor(0.27504745, shape=(), dtype=float32)
tf.Tensor(0.2738936, shape=(), dtype=float32)
tf.Tensor(0.2727481, shape=(), dtype=float32)
tf.Tensor(0.27161098, shape=(), dtype=float32)
tf.Tensor(0.27048206, shape=(), dtype=float32)
tf.Tensor(0.26936132, shape=(), dtype=float32)
tf.Tensor(0.26824862, shape=(), dtype=float32)
tf.Tensor(0.26714393, shape=(), dtype=float32)
tf.Tensor(0.26604718, shape=(), dtype=float32)
tf.Tensor(0.26495826, shape=(), dtype=float32)

这涵盖了使用 tf.distribute.Strategy API 分发自定义训练循环的最简单情况。

目前支持哪些功能？

训练 API	`MirroredStrategy`	`TPUStrategy`	`MultiWorkerMirroredStrategy`	`ParameterServerStrategy`	`CentralStorageStrategy`
自定义训练循环	受支持	受支持	受支持	实验性支持	实验性支持

示例和教程

以下是将分发策略与自定义训练循环结合使用的一些示例：

教程：使用自定义训练循环和 MirroredStrategy 进行训练。
教程：使用自定义训练循环和 MultiWorkerMirroredStrategy 进行训练。
指南：包含具有 TPUStrategy 的自定义训练循环示例。
教程：使用自定义训练循环和 ParameterServerStrategy 进行参数服务器训练。
TensorFlow Model Garden 存储库，其中包含使用各种策略实现的最新模型集合。

其他主题

本节涵盖与多个用例相关的主题。

设置 TF_CONFIG 环境变量

对于多工作节点训练，如前所述，您需要为集群中运行的每个二进制文件设置 'TF_CONFIG' 环境变量。'TF_CONFIG' 环境变量是一个 JSON 字符串，它指定了哪些任务构成了集群、它们的地址以及每个任务在集群中的角色。tensorflow/ecosystem 存储库提供了一个 Kubernetes 模板，它会为您的训练任务设置 'TF_CONFIG'。

'TF_CONFIG' 有两个组件：集群 (cluster) 和任务 (task)。

集群提供了关于训练集群的信息，它是一个包含不同类型工作（例如工作节点）的字典。在多工作节点训练中，通常有一个工作节点承担比常规工作节点更多的责任，例如保存检查点和为 TensorBoard 写入汇总文件。这样的工作节点被称为“chief”（首席）工作节点，通常索引为 0 的工作节点会被任命为首席工作节点（事实上这就是 tf.distribute.Strategy 的实现方式）。
另一方面，任务提供了关于当前任务的信息。第一个组件 cluster 在所有工作节点上是相同的，而第二个组件 task 在每个工作节点上是不同的，它指定了该工作节点的类型和索引。

'TF_CONFIG' 的一个示例是：

os.environ["TF_CONFIG"] = json.dumps({
    "cluster": {
        "worker": ["host1:port", "host2:port", "host3:port"],
        "ps": ["host4:port", "host5:port"]
    },
   "task": {"type": "worker", "index": 1}
})

此 'TF_CONFIG' 指定“集群”中有三个工作节点和两个 "ps" 任务，以及它们的主机和端口。“任务”部分指定了当前任务在“集群”中的角色——工作节点 1（第二个工作节点）。集群中的有效角色为 "chief"、"worker"、"ps" 和 "evaluator"。除非使用 tf.distribute.experimental.ParameterServerStrategy，否则不应有 "ps" 工作。

接下来是什么？

tf.distribute.Strategy 正在积极开发中。请尝试一下，并通过 GitHub issues 提供您的反馈。

使用 TensorFlow 进行分布式训练 使用收藏功能保持井然有序 根据您的偏好保存和分类内容。

概览

设置 TensorFlow

策略类型

MirroredStrategy

TPUStrategy

MultiWorkerMirroredStrategy

ParameterServerStrategy

CentralStorageStrategy

其他策略

默认策略

OneDeviceStrategy

将 tf.distribute.Strategy 与 Keras Model.fit 结合使用

目前支持哪些功能？

示例和教程

将 tf.distribute.Strategy 与自定义训练循环结合使用

目前支持哪些功能？

示例和教程

其他主题

设置 TF_CONFIG 环境变量

接下来是什么？

使用 TensorFlow 进行分布式训练