本文以兩篇官方文檔為基礎來學習TensorFlow如何進行分布式訓練，借此進入Strategy世界。

[翻譯] 使用 TensorFlow 進行分布式訓練

0x00 摘要

本文以下面兩篇官方文檔為基礎來學習TensorFlow 如何進行分布式訓練：

https://tensorflow.google.cn/guide/distributed_training（此文的信息是2.3版本之前）。

https://github.com/tensorflow/docs-l10n/blob/master/site/en-snapshot/guide/distributed_training.ipynb （此文是官方最近更新）。

本系列其他文章是：

[翻譯] TensorFlow 分布式之論文篇 "TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems"

[翻譯] TensorFlow 分布式之論文篇 "Implementation of Control Flow in TensorFlow"

[源碼解析] TensorFlow 分布式環境(1) --- 總體架構

[源碼解析] TensorFlow 分布式環境(2)---Master 靜態邏輯

[源碼解析] TensorFlow 分布式環境(3)--- Worker 靜態邏輯

[源碼解析] TensorFlow 分布式環境(4) --- WorkerCache

[源碼解析] TensorFlow 分布式環境(5) --- Session

[源碼解析] TensorFlow 分布式環境(7) --- Worker 動態邏輯

[源碼解析] TensorFlow 分布式環境(8) --- 通信機制

1. 概述

Tf.distribute.Strategy 是一個可在多個 GPU、多臺機器或 TPU 上進行分布式訓練的 TensorFlow API。使用此 API，您只需改動較少代碼就能基于現有模型和訓練代碼來實現單機多卡，多機多卡等情況的分布式訓練。

tf.distribute.Strategy 旨在實現以下目標：

• 覆蓋不同維度的用戶用例。
• 易于使用，支持多種用戶（包括研究人員和 ML 工程師等）。
• 提供開箱即用的高性能。
• 從用戶模型代碼之中解耦，這樣可以輕松切換策略。
• 支持 Custom Training Loop，Estimator，Keras。
• 支持 eager excution。

Tf.distribute.Strategy 可用于 Keras，Model.fit等高級 API，也可用來分布自定義訓練循環（以及（一般來說）使用 TensorFlow 的任何計算）。比如將模型進行構建和 model.compile() 調用封裝在 strategy.scope() 內部。

在 TensorFlow 2.x 中，您可以立即執行程序，也可以使用 tf.function在計算圖中執行。雖然 tf.distribute.Strategy 對兩種執行模式都支持，但使用 tf.function 效果最佳。建議僅將 Eager 模式用于調試，而 tf.distribute.TPUStrategy 支持此模式。盡管本指南大部分時間在討論訓練，但該 API 也可用于在不同平臺上分布評估和預測。

您在使用 tf.distribute.Strategy 只需改動少量代碼，因為我們修改了 TensorFlow 的底層組件，使其可感知策略。這些組件包括變量、層、模型、優化器、指標、摘要和檢查點。

在本指南中，我們將介紹各種類型的策略，以及如何在不同情況下使用它們。

2. 策略類型

Tf.distribute.Strategy 打算涵蓋不同軸上的許多用例。目前已支持其中的部分組合，將來還會添加其他組合。其中一些軸包括：

• 同步和異步訓練：這是通過數據并行進行分布式訓練的兩種常用方法。在同步訓練中，所有工作進程都同步地對輸入數據的不同片段進行訓練，并且會在每一步中聚合梯度。在異步訓練中，所有工作進程都獨立訓練輸入數據并異步更新變量。通常情況下，同步訓練通過全歸約（all-reduce）實現，而異步訓練通過參數服務器架構實現。
• 硬件平臺：您可能需要將訓練擴展到一臺機器上的多個 GPU 或一個網絡中的多臺機器（每臺機器擁有 0 個或多個 GPU），或擴展到 Cloud TPU 上。

要支持這些用例，有 MirroredStrategy, TPUStrategy, MultiWorkerMirroredStrategy, ParameterServerStrategy, CentralStorageStrategy 這六種策略可選。在下一部分，我們將說明當前在哪些場景中支持哪些策略。以下為快速概覽：

注：實驗性支持指不保證該 API 的兼容性。

注: 對 Estimator 的支持是有限的。其基本訓練和評估是實驗性的，高級功能（比如 scaffold）并沒有實現。如果一個用例沒有被涵蓋，您應該使用 Keras 或自定義訓練循環。不建議將 Estimator 用于新代碼，因為 Estimato r的代碼風格屬于 "v1.Session"，這很難正確編寫，而且可能會出現意想不到的行為，特別是當與 TF 2 代碼結合時。

2.1 MirroredStrategy

Tf.distribute.MirroredStrategy 支持在一臺機器的多個 GPU 上進行同步分布式訓練（單機多卡數據并行）。該策略會為每個 GPU 設備創建一個模型副本。模型中的每個變量都會在所有副本之間進行鏡像。這些變量將共同形成一個名為 MirroredVariable 的概念上的單個變量。通過應用相同的更新，這些變量保持彼此保持同步。

MirroredVariable 的同步更新只是提高了計算速度，但并不能像 CPU 并行那樣可以把內存之中的變量共享。即，顯卡并行計算只是提高速度，并不會讓用戶數據量翻倍。增加數據仍然會拋出來內存溢出錯誤。

MirroredStrategy 使用高效的全歸約（all-reduce）算法在設備之間傳遞變量更新。全歸約（all-reduce）算法通過把各個設備上的張量加起來使其聚合，并使其在每個設備上可用。這是一種非常高效的融合算法，可以顯著減少同步開銷。根據設備之間可用的通信類型，可以使用的全歸約（all-reduce）算法和實現方法有很多。默認使用 NVIDIA NCCL 作為全歸約（all-reduce）實現。您可以選擇我們提供的其他選項，也可以自己編寫。

具體如下：

圖 1 MirroredStrategy 來自 TensorFlow

以下是創建 MirroredStrategy 最簡單的方法：

mirrored_strategy = tf.distribute.MirroredStrategy()

這會創建一個 MirroredStrategy 實例，該實例使用所有對 TensorFlow 可見的 GPU，并使用 NCCL 進行跨設備通信。如果您只想使用機器上的部分 GPU，您可以這樣做：

mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/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())

2.2 TPUStrategy

您可以使用tf.distribute.experimental.TPUStrategy在張量處理單元 (TPU) 上運行 TensorFlow 訓練。TPU 是 Google 的專用 ASIC，旨在顯著加速機器學習工作負載。您可通過 Google Colab、TensorFlow Research Cloud 和 Cloud TPU 平臺進行使用。

就分布式訓練架構而言，TPUStrategy 和 MirroredStrategy 是一樣的，即實現同步分布式訓練。TPU 會在多個 TPU 核心之間實現高效的全歸約（all-reduce）和其他集合運算，并將其用于 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.experimental.TPUStrategy(cluster_resolver)

TPUClusterResolver 實例可幫助定位 TPU。在 Colab 中，您無需為其指定任何參數。

如果要將其用于 Cloud TPU，您必須：

• 在 tpu 參數中指定 TPU 資源的名稱。
• 在程序開始時顯式地初始化 TPU 系統。這是使用 TPU 進行計算前的必須步驟。初始化 TPU 系統還會清除 TPU 內存，所以為了避免丟失狀態，請務必先完成此步驟。

2.3 MultiWorkerMirroredStrategy

Tf.distribute.experimental.MultiWorkerMirroredStrategy與 MirroredStrategy 非常相似。它實現了跨多個工作進程的同步分布式訓練（多機多卡分布式版本），而每個工作進程可能有多個 GPU。與 MirroredStrategy 類似，它也會跨所有工作進程在每個設備的模型中創建所有變量的副本。

圖 2 MultiWorkerMirroredStrategy 來自 TensorFlow

它使用 CollectiveOps 作為多工作進程全歸約（all-reduce）通信方法，用于保持變量同步。集合運算是 TensorFlow 計算圖中的單個運算，它可以根據硬件、網絡拓撲和張量大小在 TensorFlow 運行期間自動選擇全歸約（all-reduce）算法。

圖 3 MultiWorkerMirroredStrategy 數據流. 來自 TensorFlow

它還實現了其他性能優化。例如，靜態優化，可以將小張量上的多個全歸約（all-reduce）轉化為大張量上較少的全歸約（all-reduce）。此外，我們還在為它設計插件架構，這樣您將來就能以插件的形式使用針對您的硬件進行了更好優化的算法。

以下是創建 MultiWorkerMirroredStrategy 最簡單的方法：

multiworker_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()

MultiWorkerMirroredStrategy 目前為您提供兩種不同的集合運算實現方法。CollectiveCommunication.RING通過將 RPC 用作通信層來實現基于環的集合，支持CPU和GPU。CollectiveCommunication.NCCL 使用 NCCL 來實現集合。CollectiveCommunication.AUTO 會將選擇推遲到運行時。集合實現的最佳選擇取決于 GPU 的數量和種類，以及集群中的網絡互連。您可以通過以下方式來指定：

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

或者

if distribution_strategy == "multi_worker_mirrored":
  return tf.distribute.experimental.MultiWorkerMirroredStrategy(
      communication=_collective_communication(all_reduce_alg))

具體如下：

def _collective_communication(all_reduce_alg):
  """Return a CollectiveCommunication based on all_reduce_alg.

  Args:
    all_reduce_alg: a string specifying which collective communication to pick,
      or None.

  Returns:
    tf.distribute.experimental.CollectiveCommunication object

  Raises:
    ValueError: if all_reduce_alg not in [None, "ring", "nccl"]
  """
  collective_communication_options = {
      None: tf.distribute.experimental.CollectiveCommunication.AUTO,
      "ring": tf.distribute.experimental.CollectiveCommunication.RING,
      "nccl": tf.distribute.experimental.CollectiveCommunication.NCCL
  }
  if all_reduce_alg not in collective_communication_options:
    raise ValueError(
        "When used with multi_worker_mirrored, valid values for "
        "all_reduce_alg are [ring, nccl].  Supplied value: {}".format(
            all_reduce_alg))
  return collective_communication_options[all_reduce_alg]

與多 GPU 訓練相比，多工作進程訓練的一個主要差異是多工作進程的設置。TF_CONFIG 環境變量是在 TensorFlow 中為作為集群一部分的每個工作進程指定集群配置的標準方法。

2.4 CentralStorageStrategy

tf.distribute.experimental.CentralStorageStrategy也執行同步訓練。變量不會被鏡像，而是統一放在 CPU 上，模型和運算會復制到所有本地 GPU（這屬于 in-graph 復制，就是一個計算圖覆蓋了多個模型副本）。如果只有一個 GPU，則所有變量和運算都將被放在該 GPU 上。這樣可以處理 embedding 無法放置在一個 GPU 之上的情況。比如下圖是：單機多個 GPU。

圖 4 CentralStorageStrategy. 來自 TensorFlow

可以通過以下代碼，創建 CentralStorageStrategy 實例：

central_storage_strategy = tf.distribute.experimental.CentralStorageStrategy()

這會創建一個 CentralStorageStrategy 實例，該實例將使用所有可見的 GPU 和 CPU。在副本上對變量的更新將先進行聚合，然后再應用于變量。

注：此策略處于 experimental 階段，我們目前正在進行改進，使其能夠用于更多場景。敬請期待 API 的未來變化。 CentralStorageStrategy 優點是 GPU 負載比較均衡，但是 CPU 和 GPU 通信代價會比較大。

2.5 ParameterServerStrategy

參數服務器訓練是一種常見的數據并行方法，可以在多臺機器上擴展訓練。一個參數服務器訓練集群由工作者和參數服務器組成。在訓練過程之中使用參數服務器來統一創建/管理變量（模型每個變量都被放在參數服務器上），變量在每個步驟中被工作者讀取和更新。計算則會被復制到所有工作進程的所有 GPU 中（注：該 V1 版本策略僅適用于 Estimator API)。

在 TensorFlow 2 中，參數服務器訓練使用了一個基于中央協調者（central coordinator-based）的架構，這通過tf.distribute.experimental.coordinator.ClusterCoordinator類來完成。

TensorFlow 2 參數服務器使用異步方式來更新，即，會在各工作節點上獨立進行變量的讀取和更新，無需采取任何同步操作。因為工作節點彼此互不依賴，因此該策略可以對工作者進行容錯處理，這樣會在使用搶占式服務器時有所助益。

在這個實現中，工作者和參數服務器運行 tf.distribution.Servers 來聽取協調者的任務。協調器負責創建資源，分配訓練任務，寫檢查點，并處理任務失敗的情況。

圖 5 ParameterServerStrategy. 來自 TensorFlow

如果要在協調器上運行，您需要使用 ParameterServerStrategy 對象來定義訓練步驟，并使用 ClusterCoordinator 將訓練步驟分派給遠程工作者。下面是創建它們的最簡單方法。

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

注：如果您使用TFConfigClusterResolver，您將需要配置 TF_CONFIG 環境變量。它類似于MultiWorkerMirroredStrategy 中的'TF_CONFIG'，但有額外的注意事項。在TensorFlow 1 中，ParameterServerStrategy 只能通過tf.compat.v1.distribution.experimental.ParameterServerStrategy 符號與 Estimator一起使用。

注：這個策略是實驗性的，它目前正在積極開發中。

2.6 其他策略

除上述策略外，還有其他兩種策略可能對使用 tf.distribute API 進行原型設計和調試有所幫助。

2.6.1 默認策略

默認策略（Default Strategy）是一種分布式策略，當作用域內沒有顯式指定分布策略時就會使用此策略。此策略會實現 tf.distribute.Strategy 接口，但只具有傳遞（pass-through）功能，不提供實際分布（distribution）。例如，strategy.run(fn) 只會調用 fn。使用該策略編寫的代碼與未使用任何策略編寫的代碼完全一樣。您可以將其視為 “無運算 no-op” 策略。

默認策略是一種單一實例，無法創建它的更多實例。可通過在任意顯式策略的作用域（與可用于在顯式策略的作用域內獲得當前策略的 API 相同）外使用 tf.distribute.get_strategy() 獲得該策略。

default_strategy = tf.distribute.get_strategy()

該策略有兩個主要用途：

• 它允許無條件編寫可感知分布的庫代碼。例如，在優化器中，我們可以執行 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.)  # reduce some values

• 與庫代碼類似，它可以使用戶程序在使用或不使用分布策略的情況下都能工作，而無需條件邏輯。以下示例代碼段展示了這一點：

if tf.config.list_physical_devices('gpu'):
  strategy = tf.distribute.MirroredStrategy()
else:  # use default strategy
  strategy = tf.distribute.get_strategy() 

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

2.6.2 OneDeviceStrategy

Tf.distribute.OneDeviceStrategy 是一種會將所有變量和計算放在單個指定設備上的策略。

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

此策略與默認策略在諸多方面存在差異。在默認策略中，與沒有任何分布策略的 TensorFlow 運行相比，變量放置邏輯保持不變。但是當使用 OneDeviceStrategy 時，在其作用域內創建的所有變量都會被顯式地放在指定設備上。此外，通過 OneDeviceStrategy.run 調用的任何函數也會被放在指定設備上。

通過該策略分布的輸入將被預提取到指定設備。而在默認策略中，則沒有輸入分布。與默認策略類似，在切換到實際分布到多個設備/機器的其他策略之前，也可以使用此策略來測試代碼。這將比默認策略更多地使用分布策略機制，但不能像使用 MirroredStrategy 或 TPUStrategy 等策略那樣充分發揮其作用。如果您想讓代碼表現地像沒有策略，請使用默認策略。

目前為止，我們已經討論了可用的不同策略以及如何將其實例化。在接下來的幾個部分中，我們將討論使用它們分布訓練的不同方法。我們將在本指南中展示簡短的代碼段，并附上可以從頭到尾運行的完整教程的鏈接。

3. 在tf.keras.Model.fit 中使用

我們已將 tf.distribute.Strategy 集成到 tf.keras。tf.keras 是用于構建和訓練模型的高級 API。將該策略集成到 tf.keras 后端以后，您可以使用 model.fit 在 Keras 訓練框架中無縫進行分布式訓練。您需要對代碼進行以下更改：

1. 創建一個合適的 tf.distribute.Strategy 實例。
2. 將 Keras 模型、優化器和指標的創建轉移到 strategy.scope 中。

我們支持所有類型的 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,))])

model.compile(loss='mse', optimizer='sgd')

在此示例中我們使用了 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)

我們在這里使用了 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)

在上述兩種情況（Dataset 或 Numpy）中，給定輸入的每個批次都被平均分到了多個副本中。例如，如果對 2 個 GPU 使用 MirroredStrategy，大小為 10 的每個批次將被均分到 2 個 GPU 中，每個 GPU 每步會接收 5 個輸入樣本。如果添加更多 GPU，每個周期的訓練速度就會更快。在添加更多加速器時通常需要增加批次大小，以便有效利用額外的計算能力。您還需要根據模型重新調整學習率。您可以使用 strategy.num_replicas_in_sync 獲得副本數量。

# Compute global batch size using 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}
learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]

目前支持的策略如下：

4. 在自定義訓練循環中使用

如您所見，在 Keras model.fit 中使用 tf.distribute.Strategy 只需改動幾行代碼。再多花點功夫，您還可以在自定義訓練循環中使用 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,))])
  optimizer = tf.keras.optimizers.SGD()

接下來，我們創建輸入數據集并調用 tf.distribute.Strategy.experimental_distribute_dataset 以根據策略來分布數據集。

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

然后，我們定義一個訓練步驟。我們將使用 tf.GradientTape 來計算梯度，并使用優化器來應用這些梯度以更新模型變量。要分布此訓練步驟，我們加入一個 train_step 函數，并將此函數和從之前創建的 dist_dataset 獲得的數據集輸入一起傳遞給 tf.distrbute.Strategy.run：

loss_object = tf.keras.losses.BinaryCrossentropy(
  from_logits=True,
  reduction=tf.keras.losses.Reduction.NONE)

def compute_loss(labels, predictions):
  per_example_loss = loss_object(labels, predictions)
  return tf.nn.compute_average_loss(per_example_loss, global_batch_size=global_batch_size)

def train_step(inputs):
  features, labels = inputs

  with tf.GradientTape() as tape:
    predictions = model(features, training=True)
    loss = compute_loss(labels, predictions)

  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)

以上代碼還需注意以下幾點：

1. 我們使用了 tf.nn.compute_average_loss 來計算損失。tf.nn.compute_average_loss 將每個樣本的損失相加，然后將總和除以 global_batch_size。這很重要，因為稍后在每個副本上計算出梯度后，會通過對它們求和使其跨副本進行聚合。
2. 我們使用了 tf.distribute.Strategy.reduceAPI 來聚合 tf.distribute.Strategy.run 返回的結果。tf.distribute.Strategy.run 會從策略中的每個本地副本返回結果。目前有多種方法使用該結果，比如可以 reduce 它們以獲得聚合值。還可以通過執行 tf.distribute.Strategy.experimental_local_results 獲得包含在結果中的值的列表，每個本地副本一個列表。
3. 當在一個分布策略作用域內調用 apply_gradients 時，它的行為會被修改。具體來說，在同步訓練期間，在將梯度應用于每個并行實例之前，它會對梯度的所有副本求和（sum-over-all-replicas）。

最后，當我們定義完訓練步驟后，就可以迭代 dist_dataset，并在循環中運行訓練：

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

在上面的示例中，我們通過迭代 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.distribute.StrategyAPI 來分布自定義訓練循環（custom training loops）最簡單的情況。

5. 其他主題

在此部分，我們將介紹與多個用例相關的主題。

5.1 設置 TF_CONFIG 環境變量

對于多工作進程訓練來說，如前所述，您需要為每個在集群中運行的二進制文件設置 TF_CONFIG 環境變量。TF_CONFIG 環境變量是一個 JSON 字符串，它指定了構成集群的任務、它們的地址，以及每個任務在集群中的角色。我們在 tensorflow/ecosystem 倉庫中提供了一個 Kubernetes 模板，可為您的訓練任務設置 TF_CONFIG。

TF_CONFIG 有兩個組件：cluster 和 task。

• cluster 會提供有關訓練集群的信息，這是一個由不同類型的作業（如工作進程）組成的字典。在多工作進程訓練中，通常會有一個工作進程除了要完成常規工作進程的工作之外，還要承擔更多責任，如保存檢查點和為 TensorBoard 編寫摘要文件。此類工作進程稱為 "chief" 工作進程，習慣上會將索引為 0 的工作進程指定為 chief 工作進程（實際上這是 tf.distribute.Strategy 的實現方式）。
• 另一方面，task 會提供有關當前任務的信息。第一個組件 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 任務，以及它們的主機和端口。"task" 部分指定當前任務在集群中的角色，即 worker 1（第二個工作進程）。集群中的有效角色是 "chief"、"worker"、"ps" 和 "evaluator"。除使用 tf.distribute.experimental.ParameterServerStrategy 時外，不應有 "ps" 作業。

0xFF 參考

使用 TensorFlow 進行分布式訓練

https://github.com/tensorflow/docs-l10n/blob/master/site/en-snapshot/guide/distributed_training.ipynb

Tensorflow上手4: 初探分布式訓練