‘Parameter server training with ParameterServerStrategy‘官方文档翻译及个人理解

0 前言

由于目前全网很少关于Parameter server training的实践，而官网分布式训练模块就这一篇文章没有翻译，读者读起来很费劲。而使用机翻经常乱七八糟，因此笔者决定在尊重官网原文的准则下自己来翻译，并偶尔给出自己的思考。
官方文档地址

1 概述

Parameter server training是一种常用的数据并行方法，用以将模型训练扩展到多台机器上。

一个参数服务器训练集群包括workers和parameter servers。变量在parameter servers上被创建，并且在训练的每一步中都被workers读取和更新。默认情况下，workers之间不同步地独立地进行读取和更新变量。这也是为什么Parameter server training被称为异步训练的原因。

在tf2中，Parameter server training被tf.distribute.ParameterServerStrategy这个API所支持，这个策略可以将训练步骤分布到扩展了数千个workers和parameter servers构成的分布式集群上。

1.1 支持的训练方法

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

1. Keras的 Model.fit API：如果你更喜欢高阶抽象和处理一个tf.keras.Model的话，这个是我们通常推荐的。
2. 自定义训练循环：如果你更喜欢自己定义训练循环的细节的话。

1.2 集群中的角色说明

不管是 Model.fit还是自定义循环，tf2中的分布式训练包括一个包含多个jobs的cluster，并且每个job可能有一个或多个tasks。

当使用Parameter server training时，我们推荐设置：

1. 一个coordinator job：（job名称为chief）
2. 多个worker jobs：（名称为worker）
3. 多个parameter server jobs：（名称为ps）

coordinator具有创建资源，调度训练任务，写入检查点，处理任务错误的作用。workers和parameter servers运行一个tf.distribute.Server实例来从coordinator监听请求。

1.3 使用 Model.fit API进行Parameter server training

这种方式要求coordinator使用tf.distribute.ParameterServerStrategy这个API对象。像不使用任何策略或者使用其他策略类似，工作流包括创建并编译模型、准备回调函数以及调用 Model.fitAPI。

1.4 自定义循环进行Parameter server training

如果使用自定义循环训练，tf.distribute.experimental.coordinator.ClusterCoordinator类是一个适用于coordinator的关键组件。

1. ClusterCoordinator类需要配合tf.distribute.StrategyAPI对象工作。
2. tf.distribute.Strategy对象需要提供集群的信息并且是用来定义一个训练step的，就像Custom training with tf.distribute.Strategy中描述的一样。
3. 然后，ClusterCoordinator类会调度这些训练的执行任务到远端的workers
4. 对于Parameter server training， ClusterCoordinator类需要配合tf.distribute.ParameterServerStrategy使用。

ClusterCoordinator提供的一个最重要的API是schedule：

1. scheduleAPI将一个tf.function加入到任务队列，并且立即返回一个将来的RemoteValue。
2. 队列中的function将会以多个后台进程的方式被调度到多个远端workers上，并且远端workers上的RemoteValue将会以异步的方式进行填充。
3. 既然schedule不要求workers进行分配和执行任务，传入的tf.function将会执行在任何一个可获得的worker上。
4. 当训练任务还没结束但是执行任务所在的worker出故障变得不可获得的话，tf.function将会在另一个可获得的worker上重新执行。
5. 由于tf.function的执行不是自动化的，一个单独的function可能会被多次执行。

除了调度function远程执行之外，ClusterCoordinator同时也有助于 在所有机器上创建数据集和当某个worker故障之后修复后进行数据集重建。

2 教程中的设置

这篇教程将会分支为 Model.fit 和自定义训练循环两种学习路径，你可以选择适合自己需求的。除了标题中含有“...进行训练”的章节，其他的这两条路径都适用。

# 终端安装端口相关组件
pip install portpicker

# 导入相关包
import multiprocessing
import os
import random
import portpicker
import tensorflow as tf

3 集群设置

像上文提到的，一个Parameter server training集群需要一个coordinator来运行训练计划，一个或多个worker和ps运行tf.distribute.Server实例，并且可能一个额外的evaluator task来执行边车评估任务（sidecar evaluation）。
设置要求如下：

1. coordinator需要知道所有的其他workers和ps的IP地址和端口，除了evaluator之外。
2. workers和ps需要知道他们应该监听哪个端口。为了简便起见，你可以传入一个完全的集群信息进去。
3. evaluator task不需要知道训练集群的设置。即使知道了，也不能尝试去链接这个训练集群。
4. workers和parameter servers应该要有分别为workers和ps的task名称。coordinator的task名称根据遗留原因一个设置为chief。

在这个教程中，你将会创建一个in-process集群来支持所有的节点都能在Colab上运行。如何设置一个真实的集群将会在后面部分中进行学习。

3.1 in-process集群

你需要事先创建多个TensorFlow servers，然后你可以连接到他们。注意：这个只是在教程中可以用，真实的训练需要运行在真实的workers和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)

这种in-process集群设置被频繁的应用到单元测试中。
另外一种进行本地测试的选项是运行多进程（我理解为不同的端口号代表不一样的机器？）：参考Multi-worker training with Keras。

4 实例化一个ParameterServerStrategy

在进行训练代码的编写之前，我们需要先实例化一个tf.distribute.ParameterServerStrategy对象。

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)

为了使用GPUs进行训练，每个worker需要设置GPU可见。ParameterServerStrategy将会让每台worker使用所有可获得的GPUs。但有一个限制就是：所有的workers需要装配同样数目的GPUs。

我有一个新的认知：只有workers需要配备GPU？！parameter servers和coordinator不硬性要求么？
回答：实践后发现，不仅不是硬性要求，反而coordinator要求不能进行使用多GPU训练，因此应设置为一个可见（或直接禁用？）

4.1 变量分片

变量分片指的是将一个较大的变量分片为多个小的变量，我们称之为shards。变量分片在访问这些shards的时候对分散网络负载十分有用，同时也对分散多个parameter servers的计算和存储能力十分有用。

为使得可进行变量分片，在构建一个ParameterServerStrategy对象时，可以传入一个变量分片器–variable_partitioner。每当变量被创建这个变量分片器便会被请求，并且期望是返回沿着变量每个维度进行分片的shards的数量。

一些开箱可用的变量分片器可供使用，比如：tf.distribute.experimental.partitioners.MinSizePartitioner。这种基于大小的变量分片器是被推荐的，可以防止对过小的变量进行分片，以及减少因过小变量而对模型训练速度产生的消极影响。

当variable_partitioner被传入时，你就已经在Strategy.scope中直接创建好了一个带有变量性质的容器类型，使得可以访问到shards。在大多数情况下，这个容器将会通过串联所有shards的方式自动地被转换成一个Tensor。结果就是，它可以作为一个通常的变量使用。另一方面，一些TensorFlow方法为这种容器类型提供了一个高效的执行力，并且在这些方法中自动串联shards的行为将会被避免。

查看tf.distribute.ParameterServerStrategy以获得更多细节。

5 Model.fit方式进行训练

Keras提供了一种方便实用的训练API–Model.fit，在底层处理训练循环，具有对训练步骤可覆写的灵活性，以及回调函数提供了保存检查点和模型summary以供TensorBoard的能力。如果使用了Model.fit，相同的训练代码只需要封装在一个策略对象内就可以被其他类型策略所使用。

5.1 输入数据

有以下三种形式输入数据：tf.data.Dataset、tf.distribute.DistributedDataset、tf.keras.utils.experimental.DatasetCreator，为了方便使用，建议选择Dataset。如果使用Dataset遭遇了存储的问题，你可能就需要使用带有dataset_fn参数的DatasetCreator，具体查看以上链接。

如果你将你的数据集转换为了tf.data.Dataset，你需要使用Dataset.shuffle和Dataset.repeat，具体使用参见下面代码例子。

1. 除非数据集被不同方式地打乱，我们假设每个worker接收到一样的数据集。因此，通过调用Dataset.shuffle可以保证数据得到一个更平均的迭代。
2. 由于worker并不是同步训练，他们可能在不同时间点结束他们数据集的进程。因此，Parameter server training进行定义epochs最简单的方式就是使用Dataset.repeat，这个方法如果调用时不指定参数就会无限重复数据集，也可以在Model.fit中指定steps_per_epoch参数（每轮训练多少步）。

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)

如果你使用的是DatasetCreator进行创建数据集，在输入设备上dataset_fn中的代码将会被请求，输入设备通常会是每台worker机器上的CPU。

5.2 模型构建和编译

现在，你将创建一个keras模型，一个简单的仅用作演示目的的顺序模型tf.keras.models.Sequential，接下来便是编译模型，会调用的组件包括：优化器（optimizer）, 度量（metrics）, 以及每次执行的步数 （steps_per_execution）。

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

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

5.3 回调函数和训练

在调用Model.fit进行实际训练之前，我们需要准备一般任务都需要的回调函数，比如：

1. tf.keras.callbacks.ModelCheckpoint：以一个特定的频率保存模型，比如每训练完一个epoch保存一次。
2. tf.keras.callbacks.BackupAndRestore：如果集群遭遇不可获得性（中断或被抢占），这个回调函数会通过备份当前模型和epoch数来提供容错能力。你可以之后通过重启这个失败来回复训练状态，并从被中断的那一个epoch继续。
3. tf.keras.callbacks.TensorBoard：周期性地将总结性文件中的模型信息写入TensorBoard工具使之能够可视化。

注意事项：出于性能考虑，当使用ParameterServerStrategy时，自定义callbacks没法进行batch层次的覆写。请修改你的自定义callbacks使之能够在epoch层次调用，并且调整steps_per_epoch到一个合适的值。另外，当ParameterServerStrategy下使用Model.fit的情况下steps_per_epoch参数是必须的

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)

5.4 ClusterCoordinator的使用方法（可选）

即使你选择的是Model.fit的训练方式，你任然可以选择实例化一个tf.distribute.coordinator.ClusterCoordinator对象来调度你想在workers上执行的其他函数。

6 自定义训练循环进行训练

使用自定义训练循环的方式提供了巨大的灵活性定义训练循环，ParameterServerStrategy下你需要使用tf.distribute.coordinator.ClusterCoordinator来调度训练步骤的执行到远程的workers上。

然后，你将创建一个模型，定义一个数据集并且定义一个step函数，就像在其他分布式策略一样。详细信息参见Custom training with tf.distribute.Strategy。

为了保证高效的数据集预提取，推荐使用ClusterCoordinator.create_per_worker_datasetAPI。另外，保证调用worker_fn中的Strategy.run来充分利用装配在workers上的所有GPUs。

6.1 设置数据集

首先，创建一个函数来创建数据集。

如果你想用Keras preprocessing layers或Tensorflow Transform layers对数据进行预处理，在dataset_fn之外以及在Strategy.scope下想创建任何keras layers一样创建这些layers。这是因为dataset_fn将会被封装进tf.function，然后在每台worker上执行来生成数据流水线。

如果你不遵循以上的程序步骤，创建layers可能会创建TensorFlow状态，并从tf.function中被提升到coordinator。因此，在worker上访问它们将会导致coordinator和worker之间反复的RPC请求，由此造成明显的减缓训练速度。Then, you will apply the transformation inside the dataset_fn via tf.data.Dataset.map. Refer to Data preprocessing in the Distributed input tutorial for more information on data preprocessing with distributed input.

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

6.2 构建模型

保证所有变量都在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.RMSprop(learning_rate=0.1)
  accuracy = tf.keras.metrics.Accuracy()

6.3 定义训练步骤

创建封装在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)

In the above training step function, calling Strategy.run and Strategy.reduce in the step_fn can support multiple GPUs per worker. If the workers have GPUs allocated, Strategy.run will distribute the datasets on multiple replicas.

6.4 将训练步骤调度到远程workers

在所有的计算被ParameterServerStrategy定义之后，你将使用tf.distribute.coordinator.ClusterCoordinator来创建资源并且将训练步骤调度到远程workers。

创建一个ClusterCoordinator并传入strategy实例对象：

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

然后，使用ClusterCoordinator.create_per_worker_datasetAPI创建一个每worker的数据集和迭代器，它将复制数据集到所有的workers上。In the per_worker_dataset_fn below, wrapping the dataset_fn into strategy.distribute_datasets_from_function is recommended to allow efficient prefetching to GPUs seamlessly.

@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来分布算力到远端workers。：

1. scheduleAPI将一个tf.function加入到任务队列，并且立即返回一个将来的RemoteValue。队列中的function将会以多个后台进程的方式被调度到多个远端workers上，并且远端workers上的RemoteValue将会以异步的方式进行填充。
2. 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." % (i, accuracy.result().numpy()))

Here is how you can fetch the result of a RemoteValue:

loss = coordinator.schedule(step_fn, args=(per_worker_iterator,))
print("Final loss is %f" % loss.fetch())

Alternatively, you can launch all steps and do something while waiting for completion:

for _ in range(total_steps):
  coordinator.schedule(step_fn, args=(per_worker_iterator,))
while not coordinator.done():
  time.sleep(10)
  # Do something like logging metrics or writing checkpoints.

For the complete training and serving workflow for this particular example, please check out this test.

6.5 更多有关数据集创建

The dataset in the above code is created using the ClusterCoordinator.create_per_worker_dataset API. It creates one dataset per worker and returns a container object. You can call the iter method on it to create a per-worker iterator. The per-worker iterator contains one iterator per worker and the corresponding slice of a worker will be substituted in the input argument of the function passed to the ClusterCoordinator.schedule method before the function is executed on a particular worker.

The ClusterCoordinator.schedule method assumes workers are equivalent and thus assumes the datasets on different workers are the same (except that they may be shuffled differently). Because of this, it is also recommended to repeat datasets, and schedule a finite number of steps instead of relying on receiving an OutOfRangeError from a dataset.

Another important note is that tf.data datasets don’t support implicit serialization and deserialization across task boundaries. So it is important to create the whole dataset inside the function passed to ClusterCoordinator.create_per_worker_dataset. The create_per_worker_dataset API can also directly take a tf.data.Dataset or tf.distribute.DistributedDataset as input.

7 评估

两种评估方式：内嵌评估（inline evaluation）和旁车评估（sidecar evaluation）。

7.1 内嵌评估（inline evaluation）

这种方法下，coordinator的作用在训练和评估之间交替，故得此名。
inline evaluation的一些优势，比如：

1. 它能支持单个task无法支持的大型评估模型和评估数据集。
2. 评估结果可以被用来下一轮训练的决策，比如，根据评估结果决策是否提前结束训练。

两种方式实现inline evaluation：直接评估和分布式评估。

1. 直接评估：适合小的模型和评估数据集，coordinator可以直接在带有评估数据集的分布式模型上运行评估任务。

eval_dataset = tf.data.Dataset.from_tensor_slices(
    feature_and_label_gen(num_examples=16)).map(
          lambda x: (
              {"features": feature_preprocess_stage(x["features"])},
              label_preprocess_stage(x["label"])
          )).batch(8)

eval_accuracy = tf.keras.metrics.Accuracy()

for batch_data, labels in eval_dataset:
  pred = model(batch_data, training=False)
  actual_pred = tf.cast(tf.greater(pred, 0.5), tf.int64)
  eval_accuracy.update_state(labels, actual_pred)

print("Evaluation accuracy: %f" % eval_accuracy.result())

2. 分布式评估：适合大模型和数据集，直接在coordinator上运行不可行。coordinator task可以通过ClusterCoordinator.schedule/ClusterCoordinator.join这两种方法将评估task分布式分摊到workers中。

with strategy.scope():
  # Define the eval metric on parameter servers.
  eval_accuracy = tf.keras.metrics.Accuracy()

@tf.function
def eval_step(iterator):
  def replica_fn(batch_data, labels):
    pred = model(batch_data, training=False)
    actual_pred = tf.cast(tf.greater(pred, 0.5), tf.int64)
    eval_accuracy.update_state(labels, actual_pred)
  batch_data, labels = next(iterator)
  strategy.run(replica_fn, args=(batch_data, labels))

def eval_dataset_fn():
  return tf.data.Dataset.from_tensor_slices(
      feature_and_label_gen(num_examples=16)).map(
          lambda x: (
              {"features": feature_preprocess_stage(x["features"])},
              label_preprocess_stage(x["label"])
          )).shuffle(16).repeat().batch(8)

per_worker_eval_dataset = coordinator.create_per_worker_dataset(eval_dataset_fn)
per_worker_eval_iterator = iter(per_worker_eval_dataset)

eval_steps_per_epoch = 2
for _ in range(eval_steps_per_epoch):
  coordinator.schedule(eval_step, args=(per_worker_eval_iterator,))
coordinator.join()
print("Evaluation accuracy: %f" % eval_accuracy.result())

7.2 旁车评估（sidecar evaluation）

创建一个通过反复读取检查点并根据最新的检查点进行评估任务的专用evaluator。

chief和 workertasks没有花费任何时间在评估上，因此，对于一个固定迭代次数的训练，整体时间要比使用其他评估方法要快。

但，它需要一个额外的evaluator task以及周期性的检查点来触发评估任务。

有两个选项来为旁车评估写评估循环：

1. 使用tf.keras.utils.SidecarEvaluatorAPI
2. 创建自定义评估循环

旁车评估仅支持一个单一任务，这意味着：

1. 保证了每个样本只评估一次。假如evaluator被占用或者重启了，从最近的检查点重启评估循环十分简单，并且在重启之前做的部分评估进程将会被丢弃。
2. 然而，将评估任务只用一个task运行，可能同时意味着整个评估循环可能会花费很长时间。
3. 如果模型的规模太大以至于evaluator的存储不在适合，那么单个旁车评估便不再适用。
   另外一个警告是：自定义评估循环可能会跳过一些检查点没有评估，因为它经常选用最新的检查点，而在一个评估epoch中，集群可能已经生成了多个检查点。

自定义评估循环提供了对细节的更多控制，比如选用哪个检查点进行评估或者提供额外的逻辑来运行评估任务。

checkpoint_dir = ...
eval_model = ...
eval_data = ...
checkpoint = tf.train.Checkpoint(model=eval_model)

for latest_checkpoint in tf.train.checkpoints_iterator(
    checkpoint_dir):
  try:
    checkpoint.restore(latest_checkpoint).expect_partial()
  except (tf.errors.OpError,) as e:
    # checkpoint may be deleted by training when it is about to read it.
    continue

  # Optionally add callbacks to write summaries.
  eval_model.evaluate(eval_data)

  # Evaluation finishes when it has evaluated the last epoch.
  if latest_checkpoint.endswith('-{}'.format(train_epochs)):
    break

8 真实环境下的集群设置

在真实的生产环境下，你将会在不同机器的不同进程中运行所有的任务。最简单的方式配置在task上集群信息就是设置"TF_CONFIG"环境变量，并使用tf.distribute.cluster_resolver.TFConfigClusterResolver来解析"TF_CONFIG"。
若你开始训练任务使用的是K8s或者其他的配置模板，这些模板就已经为你设置好了"TF_CONFIG"。

8.1 设置"TF_CONFIG"环境变量

假设有3个workers和2个parameter servers。则，worker 1 的"TF_CONFIG"可以是：（注意index）

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

evaluator 的"TF_CONFIG"可以是：（可选）

os.environ["TF_CONFIG"] = json.dumps({
    "cluster": {
        "evaluator": ["host7:port"]
    },
    "task": {"type": "evaluator", "index": 0}
})

8.2 若所有task使用了同样的设置

若所有task使用了同样的设置，你需要让你的程序在一开始就分化成不同的角色：

cluster_resolver = tf.distribute.cluster_resolver.TFConfigClusterResolver()
if cluster_resolver.task_type in ("worker", "ps"):
  # Start a TensorFlow server and wait.
elif cluster_resolver.task_type == "evaluator":
  # Run sidecar evaluation
else:
  # Run the coordinator.

下面的代码是在workers和ps上运行的，启动一个TensorFlow server并且等待：

# 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"

server = tf.distribute.Server(
    cluster_resolver.cluster_spec(),
    job_name=cluster_resolver.task_type,
    task_index=cluster_resolver.task_id,
    protocol=cluster_resolver.rpc_layer or "grpc",
    start=True)
server.join()

9 处理task错误

9.1 worker的错误

Both the tf.distribute.coordinator.ClusterCoordinator custom training loop and Model.fit approaches provide built-in fault tolerance for worker failure. Upon worker recovery, the ClusterCoordinator invokes dataset re-creation on the workers.

9.2 ps或coordinator的错误

然而，当coordinator看到了ps出错，它会立即引发一个UnavailableError 或 AbortedError。这种情况下，你可以重启coordinator
coordinator也可能会出错，因此，一些特定的工具被推荐来防止丢失训练进程：

1. 对于Model.fit，你应该使用一个BackupAndRestore的回调函数，可以自动保存并用于重启。
2. 对于自定义循环，你应该周期性地检查模型变量并根据检查点加载模型变量。如果优化器被设置了检查点，可以根据optimizer.iterations大约推断出训练进程。

checkpoint_manager = tf.train.CheckpointManager(
    tf.train.Checkpoint(model=model, optimizer=optimizer),
    checkpoint_dir,
    max_to_keep=3)
if checkpoint_manager.latest_checkpoint:
  checkpoint = checkpoint_manager.checkpoint
  checkpoint.restore(
      checkpoint_manager.latest_checkpoint).assert_existing_objects_matched()

global_steps = int(optimizer.iterations.numpy())
starting_epoch = global_steps // steps_per_epoch

for _ in range(starting_epoch, num_epochs):
  for _ in range(steps_per_epoch):
    coordinator.schedule(step_fn, args=(per_worker_iterator,))
  coordinator.join()
  checkpoint_manager.save()

9.3 获取RemoteValue

如若函数执行成功，则保证能获取到RemoteValue。因为函数一执行完返回值就被拷贝到coordinator。如果在复制的过程中发生了worker错误，则会在另外一个可获得的worker上重新尝试。因此，如果你想为性能进行优化，你可以调度functions但不返回值。

10 性能提升

当你使用tf.distribute.ParameterServerStrategy 和 tf.distribute.coordinator.ClusterCoordinator进行训练时，有很多因素可能会导致你面临性能问题。

一个常见的原因就是parameter servers没有负载均衡并且一些大负载的parameter servers已经到达了能力上限。也可能有多种根源，一些简单的方法可以缓解这个问题：

1. 模型分片
2. 如果可以的话，避免创建一个所有parameter servers在一步中同时需要的热点变量。
3. 在将large vocabularies传入keras 预处理层之前先打乱它。

另外一个性能问题的原因是coordinator。schedule/join的执行是基于Python的，因此可能超过了多线程限制。并且，coordinator和workers之间的延迟可能很大。应对方法：

1. 对于Model.fit，你可以在Model.compile中设置steps_per_execution的值为一个大于1的数
2. 对于自定义循环，你可以打包多步到一个tf.function:

steps_per_invocation = 10

@tf.function
def step_fn(iterator):
  for _ in range(steps_per_invocation):
    features, labels = next(iterator)
    def replica_fn(features, labels):
      ...

    strategy.run(replica_fn, args=(features, labels))

11 限制和总结

11.1 ParameterServerStrategy总体总结

1. os.environment["grpc_fail_fast"]="use_caller"是包括coordinator在内的所有task所必须的，来保证容错能力。
2. 不支持同步训练
3. 对于自定义循环，你可以打包多步到一个tf.function来优化性能。
4. It is not supported to load a saved_model via tf.saved_model.load containing sharded variables. Note loading such a saved_model using TensorFlow Serving is expected to work
5. 若未重启coordinator task，不支持从ps错误中恢复。
6. 创建变量应包含在Strategy.scope中，否则，资源将可能会被配置在coordinator。

11.2 Model.fit特别总结

1. steps_per_epoch是必须的。选择一个合适的值来保证每一轮中各步之间一个合适的间隔。
2. 当使用ParameterServerStrategy时，自定义callbacks没法进行batch层次的覆写。请修改你的自定义callbacks使之能够在epoch层次调用，并且调整steps_per_epoch到一个合适的值。
3. 由于某些原因，不像其他策略，ps strategy中进度条和度量将会仅在轮次边界记录下来。
4. run_eagerly不被支持。

11.3 自定义循环总结

1. ClusterCoordinator.schedule doesn’t support visitation guarantees for a dataset.
2. When ClusterCoordinator.create_per_worker_dataset is used with a callable as input, the whole dataset must be created inside the function passed to it.
3. tf.data.Options is ignored in a dataset created by ClusterCoordinator.create_per_worker_dataset.