[翻譯] 使用 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 分散式之論文篇 "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 訓練框架中無縫進行分散式訓練。您需要對程式碼進行以下更改：
- 建立一個合適的 tf.distribute.Strategy 例項。
- 將 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)
以上程式碼還需注意以下幾點：
- 我們使用了 tf.nn.compute_average_loss 來計算損失。tf.nn.compute_average_loss 將每個樣本的損失相加，然後將總和除以 global_batch_size。這很重要，因為稍後在每個副本上計算出梯度後，會通過對它們求和使其跨副本進行聚合。
- 我們使用了 tf.distribute.Strategy.reduceAPI 來聚合 tf.distribute.Strategy.run 返回的結果。tf.distribute.Strategy.run 會從策略中的每個本地副本返回結果。目前有多種方法使用該結果，比如可以 reduce 它們以獲得聚合值。還可以通過執行 tf.distribute.Strategy.experimental_local_results 獲得包含在結果中的值的列表，每個本地副本一個列表。
- 當在一個分佈策略作用域內呼叫 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" 作業。