使用 ParameterServerStrategy 进行参数服务器训练

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

概述

参数服务器训练是一种常见的数据并行方法，用于在多台机器上扩展模型训练。

参数服务器训练集群由工作进程和参数服务器组成。变量在参数服务器上创建，并在每个步骤中由工作进程读取和更新。默认情况下，工作进程会独立读取和更新这些变量，而不会彼此同步。因此，参数服务器式训练有时也称为异步训练。

在 TensorFlow 2 中，参数服务器训练由 tf.distribute.ParameterServerStrategy 类提供支持，该类会将训练步骤分发到一个集群，该集群可扩展到数千个工作进程（伴随着参数服务器）。

支持的训练方法

支持的训练方法主要有两种：

Keras Model.fit API：如果您更喜欢高级抽象和训练处理。如果您正在训练 tf.keras.Model，通常建议使用此方法。
自定义训练循环：如果您更喜欢定义训练循环的详细信息（有关详细信息，请参阅自定义训练、从头开始编写训练循环和使用 Keras 和 MultiWorkerMirroredStrategy 自定义训练循环的指南）。

具有作业和任务的集群

无论选择哪种 API（Model.fit 或自定义训练循环），TensorFlow 2 中的分布式训练都涉及一个包含多个 'jobs' 的 'cluster'，每个作业可能有一个或多个 'tasks'。

在使用参数服务器训练时，建议具有：

一个协调器作业（作业名称为 chief）
多个工作进程作业（作业名称为 worker）
多个参数服务器作业（作业名称为 ps）

协调器会创建资源、调度训练任务、编写检查点并处理任务失败。工作进程和参数服务器会运行 tf.distribute.Server 实例来监听来自协调器的请求。

使用 `Model.fit` API 进行参数服务器训练

使用 Model.fit API 训练参数服务器需要协调器使用 tf.distribute.ParameterServerStrategy 对象。与没有策略或具有其他策略的 Model.fit 用法类似，工作流包括创建和编译模型、准备回调和调用 Model.fit。

使用自定义训练循环进行参数服务器训练

对于自定义训练循环，tf.distribute.coordinator.ClusterCoordinator 类是用于协调器的关键组件。

ClusterCoordinator 类需要与 tf.distribute.ParameterServerStrategy 对象结合使用。
此 tf.distribute.Strategy 对象用于提供集群的信息并用于定义训练步骤，如使用 tf.distribute.Strategy 进行自定义训练中所示。
之后，ClusterCoordinator 对象会将这些训练步骤的执行分派给远程工作进程。

ClusterCoordinator 对象提供的最重要的 API 是 schedule：

schedule API 会将 tf.function 排入队列并立即返回一个类似未来的 RemoteValue 。
排队的函数将被分派给后台线程中的远程工作进程，并且它们的 RemoteValue 将被异步填充。
由于 schedule 不需要分配工作进程，因此传入的 tf.function 可以在任何可用的工作进程上执行。
如果执行它的工作进程在完成之前变得不可用，则该函数将在另一个可用的工作进程上重试。
由于上述事实以及函数执行非原子方式的事实，单个函数调用可能会执行多次。

除了调度远程函数外，ClusterCoordinator 还会帮助在所有工作进程上创建数据集，并在工作进程从故障中恢复后重建这些数据集。

教程设置

本教程将分支到 Model.fit 和自定义训练循环路径，您可以选择适合您需求的路径。“使用 X 进行训练”以外的部分两种路径均适用。

pip install portpicker

import multiprocessing
import os
import random
import portpicker
import tensorflow as tf

2023-11-07 23:32:09.863884: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2023-11-07 23:32:09.863934: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2023-11-07 23:32:09.865658: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

集群设置

如上所述，参数服务器训练集群需要一个运行您的训练程序的协调器任务、一个或多个工作进程和运行 TensorFlow 服务器的参数服务器任务 tf.distribute.Server，以及可能运行边车评估的附加评估任务（请参阅下面的边车评估部分）。设置它们的要求是：

协调器任务需要知道除评估器之外的所有其他 TensorFlow 服务器的地址和端口。
工作进程和参数服务器需要知道他们需要监听哪个端口。为简单起见，在这些任务上创建 TensorFlow 服务器时，通常可以传入完整的集群信息。
评估器任务不必知道训练集群的设置。如果知道，则它不应该尝试连接到训练集群。
工作进程和参数服务器的任务类型应该分别为 "worker" 和 "ps"。出于遗留原因，协调器应使用 "chief" 作为任务类型。

在本教程中，您将创建一个进程内集群，以便整个参数服务器训练可以在 Colab 中运行。您将在后面的部分中学习如何设置真正的集群。

进程内集群

您将首先创建多个 TensorFlow 服务器，稍后再连接它们。请注意，这仅用于本教程的演示目的，在实际训练中，服务器将在 "worker" 和 "ps" 机器上启动。

def create_in_process_cluster(num_workers, num_ps):
  """Creates and starts local servers and returns the cluster_resolver."""
  worker_ports = [portpicker.pick_unused_port() for _ in range(num_workers)]
  ps_ports = [portpicker.pick_unused_port() for _ in range(num_ps)]

  cluster_dict = {}
  cluster_dict["worker"] = ["localhost:%s" % port for port in worker_ports]
  if num_ps > 0:
    cluster_dict["ps"] = ["localhost:%s" % port for port in ps_ports]

  cluster_spec = tf.train.ClusterSpec(cluster_dict)

  # Workers need some inter_ops threads to work properly.
  worker_config = tf.compat.v1.ConfigProto()
  if multiprocessing.cpu_count() < num_workers + 1:
    worker_config.inter_op_parallelism_threads = num_workers + 1

  for i in range(num_workers):
    tf.distribute.Server(
        cluster_spec,
        job_name="worker",
        task_index=i,
        config=worker_config,
        protocol="grpc")

  for i in range(num_ps):
    tf.distribute.Server(
        cluster_spec,
        job_name="ps",
        task_index=i,
        protocol="grpc")

  cluster_resolver = tf.distribute.cluster_resolver.SimpleClusterResolver(
      cluster_spec, rpc_layer="grpc")
  return cluster_resolver

# Set the environment variable to allow reporting worker and ps failure to the
# coordinator. This is a workaround and won't be necessary in the future.
os.environ["GRPC_FAIL_FAST"] = "use_caller"

NUM_WORKERS = 3
NUM_PS = 2
cluster_resolver = create_in_process_cluster(NUM_WORKERS, NUM_PS)

进程内集群设置经常用于单元测试，例如这里。

本地测试的另一个选择是在本地机器上启动进程，请查看使用 Keras 进行多工作进程训练以获取这种方式的示例。

实例化 ParameterServerStrategy

在深入了解训练代码之前，我们来实例化一个 tf.distribute.ParameterServerStrategy 对象。请注意，无论您是使用 Model.fit 还是自定义训练循环，都需要这样做。variable_partitioner 参数将在变量分片部分进行解释。

variable_partitioner = (
    tf.distribute.experimental.partitioners.MinSizePartitioner(
        min_shard_bytes=(256 << 10),
        max_shards=NUM_PS))

strategy = tf.distribute.ParameterServerStrategy(
    cluster_resolver,
    variable_partitioner=variable_partitioner)

INFO:tensorflow:`tf.distribute.experimental.ParameterServerStrategy` is initialized with cluster_spec: ClusterSpec({'ps': ['localhost:38883', 'localhost:43499'], 'worker': ['localhost:38671', 'localhost:45527', 'localhost:38193']})
INFO:tensorflow:ParameterServerStrategyV2 is now connecting to cluster with cluster_spec: ClusterSpec({'ps': ['localhost:38883', 'localhost:43499'], 'worker': ['localhost:38671', 'localhost:45527', 'localhost:38193']})
INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ['/job:chief/replica:0/task:0/device:GPU:0', '/job:chief/replica:0/task:0/device:GPU:1', '/job:chief/replica:0/task:0/device:GPU:2', '/job:chief/replica:0/task:0/device:GPU:3'], variable_device = '/device:CPU:0'
INFO:tensorflow:ParameterServerStrategy (CentralStorageStrategy if you are using a single machine) with compute_devices = ['/job:chief/replica:0/task:0/device:GPU:0', '/job:chief/replica:0/task:0/device:GPU:1', '/job:chief/replica:0/task:0/device:GPU:2', '/job:chief/replica:0/task:0/device:GPU:3'], variable_device = '/device:CPU:0'
INFO:tensorflow:Number of GPUs on workers: 4
INFO:tensorflow:Number of GPUs on workers: 4

为了使用 GPU 进行训练，请分配对每个工作进程可见的 GPU。ParameterServerStrategy 将使用每个工作进程上的所有可用 GPU，但限制是所有工作进程应该有相同数量的可用 GPU。

变量分片

变量分片是指将一个变量拆分为多个较小的变量，这些变量称为分片。在访问这些分片时，变量分片可能有助于分配网络负载。这对在多个参数服务器之间分布计算和存储普通变量也很有用，例如，当使用可能不适合单个机器内存的非常大的嵌入时。

要启用变量分片，您可以在构造 ParameterServerStrategy 对象时传入 variable_partitioner。每次创建变量时都会调用 variable_partitioner，它预计会返回该变量每个维度上的分片数。提供了一些开箱即用的 variable_partitioner，例如 tf.distribute.experimental.partitioners.MinSizePartitioner。建议使用基于大小的分区程序（如 tf.distribute.experimental.partitioners.MinSizePartitioner）以避免对小变量进行分区，否则可能会对模型训练速度产生负面影响。

当传入 variable_partitioner，并且您直接在 Strategy.scope 下创建变量时，该变量将成为具有 variables 属性的容器类型，该属性提供对分片列表的访问。在大多数情况下，此容器将通过连接所有分片自动转换为张量。因此，它可以用作普通变量。另一方面，一些 TensorFlow 方法（如 tf.nn.embedding_lookup）为这种容器类型提供了有效的实现，并且在这些方法中将避免自动连接。

有关详细信息，请参阅 tf.distribute.ParameterServerStrategy 的 API 文档。

使用 `Model.fit` 进行训练

Keras 通过 Model.fit 提供了一个易于使用的训练 API，它在后台处理训练循环，具有可覆盖的 train_step 的灵活性，以及为 TensorBoard 提供检查点保存或摘要保存等功能的回调。使用 Model.fit，只需简单地交换策略对象，即可将相同的训练代码用于其他策略。

输入数据

带有 tf.distribute.ParameterServerStrategy 的 Keras Model.fit 可以采用以下形式的输入数据：tf.data.Dataset、tf.distribute.DistributedDataset、tf.keras.utils.experimental.DatasetCreator，或为了使用方便推荐的 Dataset 选项。但是，如果您在使用 Dataset 时遇到内存问题，可能需要将 DatasetCreator 与可调用的 dataset_fn 参数一起使用（有关详细信息，请参阅 tf.keras.utils.experimental.DatasetCreator API 文档）。

如果将数据集转换为 tf.data.Dataset，则应使用 Dataset.shuffle 和 Dataset.repeat，如下面的代码示例所示。

带有参数服务器训练的 Keras Model.fit 假设每个工作进程接收相同的数据集（除非以不同的方式乱序）。因此，您可以通过调用 Dataset.shuffle，确保对数据进行更均匀的迭代。
由于工作进程不同步，它们可能会在不同的时间完成对数据集的处理。因此，使用参数服务器训练定义周期的最简单方式是使用 Dataset.repeat（在不带参数的情况下无限重复数据集）并在 Model.fit 调用中指定 steps_per_epoch 参数。

有关 shuffle 和 repeat 的更多详细信息，请参阅 tf.data 指南中的“训练工作流”部分。

global_batch_size = 64

x = tf.random.uniform((10, 10))
y = tf.random.uniform((10,))

dataset = tf.data.Dataset.from_tensor_slices((x, y)).shuffle(10).repeat()
dataset = dataset.batch(global_batch_size)
dataset = dataset.prefetch(2)

如果您改为使用 tf.keras.utils.experimental.DatasetCreator 创建数据集，dataset_fn 中的代码将在每台工作进程机器的输入设备（通常是 CPU）上调用。

模型构建和编译

接下来，您将创建 tf.keras.Model，这是一个用于演示目的的简单 tf.keras.models.Sequential 模型。随后调用 Model.compile 以合并组件，例如优化器、指标和其他参数（例如steps_per_execution）：

with strategy.scope():
  model = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])

  model.compile(tf.keras.optimizers.legacy.SGD(), loss="mse", steps_per_execution=10)

回调和训练

在调用 Keras Model.fit 进行实际训练之前，请为常见任务准备任何所需的回调，例如：

tf.keras.callbacks.ModelCheckpoint：以特定频率保存模型，例如在每个周期之后。
tf.keras.callbacks.BackupAndRestore：如果集群遇到不可用（例如中止或抢占），则通过备份模型和当前周期数来提供容错。然后，您可以在作业失败后重新启动时恢复训练状态，并从中断的周期开始继续训练。
tf.keras.callbacks.TensorBoard：定期将模型日志写入可在 TensorBoard 工具中可视化的摘要文件中。

注：出于性能考虑，自定义回调在与 ParameterServerStrategy 一起使用时不能覆盖批处理级别的回调。请修改您的自定义回调以进行周期级别调用，并将 steps_per_epoch 调整为合适的值。此外，当与 ParameterServerStrategy 一起使用时，steps_per_epoch 是 Model.fit 的必需参数。

working_dir = "/tmp/my_working_dir"
log_dir = os.path.join(working_dir, "log")
ckpt_filepath = os.path.join(working_dir, "ckpt")
backup_dir = os.path.join(working_dir, "backup")

callbacks = [
    tf.keras.callbacks.TensorBoard(log_dir=log_dir),
    tf.keras.callbacks.ModelCheckpoint(filepath=ckpt_filepath),
    tf.keras.callbacks.BackupAndRestore(backup_dir=backup_dir),
]

model.fit(dataset, epochs=5, steps_per_epoch=20, callbacks=callbacks)

/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/data/ops/dataset_ops.py:462: UserWarning: To make it possible to preserve tf.data options across serialization boundaries, their implementation has moved to be part of the TensorFlow graph. As a consequence, the options value is in general no longer known at graph construction time. Invoking this method in graph mode retains the legacy behavior of the original implementation, but note that the returned value might not reflect the actual value of the options.
  warnings.warn("To make it possible to preserve tf.data options across "
Epoch 1/5
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Reduce to /device:CPU:0 then broadcast to ('/replica:0/device:CPU:0',).
INFO:tensorflow:Waiting for all global closures to be finished.
INFO:tensorflow:Waiting for all global closures to be finished.
INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets
INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets
20/20 - 4s - loss: 0.9051 - 4s/epoch - 216ms/step
Epoch 2/5
INFO:tensorflow:Waiting for all global closures to be finished.
INFO:tensorflow:Waiting for all global closures to be finished.
INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets
INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets
20/20 - 1s - loss: 0.7227 - 1s/epoch - 60ms/step
Epoch 3/5
INFO:tensorflow:Waiting for all global closures to be finished.
INFO:tensorflow:Waiting for all global closures to be finished.
WARNING:tensorflow:5 out of the last 5 calls to <function MultiDeviceSaver.save.<locals>.tf_function_save at 0x7f475878cb80> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
WARNING:tensorflow:5 out of the last 5 calls to <function MultiDeviceSaver.save.<locals>.tf_function_save at 0x7f475878cb80> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets
INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets
WARNING:tensorflow:6 out of the last 6 calls to <function MultiDeviceSaver.save.<locals>.tf_function_save at 0x7f47585c0c10> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
WARNING:tensorflow:6 out of the last 6 calls to <function MultiDeviceSaver.save.<locals>.tf_function_save at 0x7f47585c0c10> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.
20/20 - 1s - loss: 0.5900 - 578ms/epoch - 29ms/step
Epoch 4/5
INFO:tensorflow:Waiting for all global closures to be finished.
INFO:tensorflow:Waiting for all global closures to be finished.
INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets
INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets
20/20 - 1s - loss: 0.4859 - 561ms/epoch - 28ms/step
Epoch 5/5
INFO:tensorflow:Waiting for all global closures to be finished.
INFO:tensorflow:Waiting for all global closures to be finished.
INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets
INFO:tensorflow:Assets written to: /tmp/my_working_dir/ckpt/assets
20/20 - 1s - loss: 0.4067 - 568ms/epoch - 28ms/step
<keras.src.callbacks.History at 0x7f49381e95e0>

直接使用 `ClusterCoordinator`（可选）

即使您选择 Model.fit 训练路径，您也可以选择实例化一个 tf.distribute.coordinator.ClusterCoordinator 对象来调度您希望在工作进程上执行的其他函数。有关更多详细信息和示例，请参阅使用自定义训练循环进行训练部分。

使用自定义训练循环进行训练

使用带有 tf.distribute.Strategy 的自定义训练循环为定义训练循环提供了极大的灵活性。通过上面定义的 ParameterServerStrategy（作为 strategy），您将使用 tf.distribute.coordinator.ClusterCoordinator 将训练步骤的执行分派给远程工作进程。

然后，您将创建一个模型、定义一个数据集并定义一个步骤函数，就像您在训练循环中使用其他 tf.distribute.Strategy 一样。您可以在使用 tf.distribute.Strategy 进行自定义训练教程中找到更多详细信息。

为确保高效的数据集预提取，请使用以下向远程工作进程分派训练步骤部分中提到的推荐分布式数据集创建 API。此外，请确保在 worker_fn 中调用 Strategy.run 以充分利用分配给工作进程的 GPU。对于使用或不使用 GPU 的训练，其余步骤相同。

我们来按以下步骤创建这些组件：

设置数据

首先，编写一个创建数据集的函数。

如果您想使用 Keras 预处理层或 Tensorflow Transform 层对数据进行预处理，请在 dataset_fn 之外和 Strategy.scope 下创建这些层，就像您对任何其他 Keras 层所做的那样。这是因为 dataset_fn 将被封装到一个 tf.function 中，然后在每个工作进程上执行以生成数据流水线。

如果您不遵循上述过程，创建层可能会创建 Tensorflow 状态，这些状态将从 tf.function 提升到协调器。因此，在工作进程上访问它们会导致协调器和工作进程之间重复的 RPC 调用，并导致显著的速度下降。。

将这些层放在 Strategy.scope 下将改为在所有工作进程上创建它们。然后，您将通过 tf.data.Dataset.map 在 dataset_fn 中应用转换。有关使用分布式输入进行数据预处理的更多信息，请参阅分布式输入教程中的数据预处理。

feature_vocab = [
    "avenger", "ironman", "batman", "hulk", "spiderman", "kingkong", "wonder_woman"
]
label_vocab = ["yes", "no"]

with strategy.scope():
  feature_lookup_layer = tf.keras.layers.StringLookup(
      vocabulary=feature_vocab,
      mask_token=None)
  label_lookup_layer = tf.keras.layers.StringLookup(
      vocabulary=label_vocab,
      num_oov_indices=0,
      mask_token=None)

  raw_feature_input = tf.keras.layers.Input(
      shape=(3,),
      dtype=tf.string,
      name="feature")
  feature_id_input = feature_lookup_layer(raw_feature_input)
  feature_preprocess_stage = tf.keras.Model(
      {"features": raw_feature_input},
      feature_id_input)

  raw_label_input = tf.keras.layers.Input(
      shape=(1,),
      dtype=tf.string,
      name="label")
  label_id_input = label_lookup_layer(raw_label_input)

  label_preprocess_stage = tf.keras.Model(
      {"label": raw_label_input},
      label_id_input)

在数据集中生成演练样本：

def feature_and_label_gen(num_examples=200):
  examples = {"features": [], "label": []}
  for _ in range(num_examples):
    features = random.sample(feature_vocab, 3)
    label = ["yes"] if "avenger" in features else ["no"]
    examples["features"].append(features)
    examples["label"].append(label)
  return examples

examples = feature_and_label_gen()

然后，创建封装在 dataset_fn 中的训练数据集：

def dataset_fn(_):
  raw_dataset = tf.data.Dataset.from_tensor_slices(examples)

  train_dataset = raw_dataset.map(
      lambda x: (
          {"features": feature_preprocess_stage(x["features"])},
          label_preprocess_stage(x["label"])
      )).shuffle(200).batch(32).repeat()
  return train_dataset

构建模型

接下来，创建模型和其他对象。确保在 Strategy.scope 下创建所有变量。

# These variables created under the `Strategy.scope` will be placed on parameter
# servers in a round-robin fashion.
with strategy.scope():
  # Create the model. The input needs to be compatible with Keras processing layers.
  model_input = tf.keras.layers.Input(
      shape=(3,), dtype=tf.int64, name="model_input")

  emb_layer = tf.keras.layers.Embedding(
      input_dim=len(feature_lookup_layer.get_vocabulary()), output_dim=16384)
  emb_output = tf.reduce_mean(emb_layer(model_input), axis=1)
  dense_output = tf.keras.layers.Dense(units=1, activation="sigmoid")(emb_output)
  model = tf.keras.Model({"features": model_input}, dense_output)

  optimizer = tf.keras.optimizers.legacy.RMSprop(learning_rate=0.1)
  accuracy = tf.keras.metrics.Accuracy()

我们来确认一下，使用 FixedShardsPartitioner 将所有变量拆分成两个分片，并且每个分片都分配给了不同的参数服务器：

assert len(emb_layer.weights) == 2
assert emb_layer.weights[0].shape == (4, 16384)
assert emb_layer.weights[1].shape == (4, 16384)

print(emb_layer.weights[0].device)
print(emb_layer.weights[1].device)

/job:ps/replica:0/task:1/device:CPU:0
/job:ps/replica:0/task:0/device:CPU:0

定义训练步骤

第三步，创建封装在 tf.function 中的训练步骤：

@tf.function
def step_fn(iterator):

  def replica_fn(batch_data, labels):
    with tf.GradientTape() as tape:
      pred = model(batch_data, training=True)
      per_example_loss = tf.keras.losses.BinaryCrossentropy(
          reduction=tf.keras.losses.Reduction.NONE)(labels, pred)
      loss = tf.nn.compute_average_loss(per_example_loss)
      gradients = tape.gradient(loss, model.trainable_variables)

    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    actual_pred = tf.cast(tf.greater(pred, 0.5), tf.int64)
    accuracy.update_state(labels, actual_pred)
    return loss

  batch_data, labels = next(iterator)
  losses = strategy.run(replica_fn, args=(batch_data, labels))
  return strategy.reduce(tf.distribute.ReduceOp.SUM, losses, axis=None)

在上面的训练步骤函数中，在 step_fn 中调用Strategy.run 和 Strategy.reduce 可以支持每个工作进程使用多个 GPU。如果工作进程分配了 GPU，则 Strategy.run 会将数据集分布在多个副本上。

向远程工作进程分派训练步骤

在 ParameterServerStrategy 定义了所有计算之后，您将使用 tf.distribute.coordinator.ClusterCoordinator 类创建资源并将训练步骤分发给远程工作进程。

我们先创建一个 ClusterCoordinator 对象，传入策略对象：

coordinator = tf.distribute.coordinator.ClusterCoordinator(strategy)

然后，使用 ClusterCoordinator.create_per_worker_dataset API 创建每个工作进程的数据集和迭代器，这会将数据集复制到所有工作进程。在下面的 per_worker_dataset_fn 中，建议将 dataset_fn 封装到 strategy.distribute_datasets_from_function 中，以便无缝且高效地预提取到 GPU。

@tf.function
def per_worker_dataset_fn():
  return strategy.distribute_datasets_from_function(dataset_fn)

per_worker_dataset = coordinator.create_per_worker_dataset(per_worker_dataset_fn)
per_worker_iterator = iter(per_worker_dataset)

最后一步是使用 ClusterCoordinator.schedule 将计算分配给远程工作进程：

schedule 方法会将一个 tf.function 排入队列，并立即返回一个类似未来的 RemoteValue。排队的函数将被分派给后台线程中的远程工作进程，并且 RemoteValue 将被异步填充。
join 方法 (ClusterCoordinator.join) 可用于等待所有调度函数执行完毕。

num_epochs = 4
steps_per_epoch = 5
for i in range(num_epochs):
  accuracy.reset_states()
  for _ in range(steps_per_epoch):
    coordinator.schedule(step_fn, args=(per_worker_iterator,))
  # Wait at epoch boundaries.
  coordinator.join()
  print("Finished epoch %d, accuracy is %f." % (

使用 ParameterServerStrategy 进行参数服务器训练 使用集合让一切井井有条 根据您的偏好保存内容并对其进行分类。

概述