TensorFlow 分散式訓練

在 TensorFlow.org 上檢視

在 Google Colab 中執行

在 GitHub 上檢視原始碼

下載筆記本

總覽

tf.distribute.Strategy 是一種 TensorFlow API，可將訓練分散到多個 GPU、多部機器或 TPU。使用此 API，您只需進行最少的程式碼變更，即可分散現有的模型和訓練程式碼。

tf.distribute.Strategy 的設計宗旨是達成以下主要目標

• 易於使用，並支援多個使用者群體，包括研究人員、機器學習工程師等。
• 提供良好的開箱即用效能。
• 輕鬆切換策略。

您可以使用 tf.distribute.Strategy，搭配 Keras Model.fit 等高階 API，以及自訂訓練迴圈 (一般而言，任何使用 TensorFlow 的運算)，來分散訓練。

在 TensorFlow 2.x 中，您可以急切地執行程式，或使用 tf.function 在圖中執行。tf.distribute.Strategy 旨在支援這兩種執行模式，但搭配 tf.function 時效果最佳。建議僅將急切模式用於偵錯目的，且 tf.distribute.TPUStrategy 不支援此模式。雖然本指南的重點是訓練，但此 API 也可用於分散不同平台上的評估和預測作業。

您可以使用 tf.distribute.Strategy，對程式碼進行極少的變更，因為 TensorFlow 的基礎元件已變更為可感知策略。這包括變數、層、模型、最佳化工具、指標、摘要和檢查點。

在本指南中，您將瞭解各種策略類型，以及如何在不同情況下使用這些策略。若要瞭解如何偵錯效能問題，請參閱最佳化 TensorFlow GPU 效能指南。

設定 TensorFlow

import tensorflow as tf

策略類型

tf.distribute.Strategy 旨在涵蓋沿著不同軸的多種使用案例。目前支援其中一些組合，其他組合將在未來新增。其中一些軸為

• 同步與非同步訓練：這是分散式訓練搭配資料平行處理的兩種常見方式。在同步訓練中，所有工作站同步訓練不同輸入資料切片，並在每個步驟彙總梯度。在非同步訓練中，所有工作站都獨立訓練輸入資料，並以非同步方式更新變數。一般而言，同步訓練透過 all-reduce 支援，非同步訓練則透過參數伺服器架構支援。
• 硬體平台：您可能會想要將訓練擴展到一部機器上的多個 GPU，或網路中的多部機器 (每部機器有 0 個或多個 GPU)，或 Cloud 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 Collective Communication Library (NCCL) 作為 all-reduce 實作。您可以從其他幾個選項中選擇，或編寫自己的選項。

以下是建立 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())

TPUStrategy

tf.distribute.TPUStrategy 可讓您在 張量處理單元 (TPU) 上執行 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()

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)

與多 GPU 訓練相比，讓多工作站訓練運作的其中一個主要差異是多工作站設定。'TF_CONFIG' 環境變數是 TensorFlow 中指定叢集組態給叢集中每個工作站的標準方式。請在本文件的設定 TF_CONFIG 一節中深入瞭解。

如需 MultiWorkerMirroredStrategy 的更多詳細資訊，請考慮以下教學課程

• 搭配 Keras Model.fit 的多工作站訓練
• 搭配自訂訓練迴圈的多工作站訓練

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()

這會建立 CentralStorageStrategy 執行個體，它將使用所有可見的 GPU 和 CPU。複本上變數的更新將在套用至變數之前彙總。

其他策略

除了上述策略之外，還有其他兩種策略可能適用於在使用 tf.distribute API 時進行原型設計和偵錯。

預設策略

預設策略是一種發布策略，當範圍內沒有明確的發布策略時，就會出現這種策略。它實作 tf.distribute.Strategy 介面，但它是直通，而且不提供實際發布。例如，Strategy.run(fn) 只會呼叫 fn。使用此策略編寫的程式碼的行為應與未編寫任何策略的程式碼完全相同。您可以將其視為「無運算」策略。

預設策略是單例模式，而且無法建立更多執行個體。它可以使用任何明確策略範圍之外的 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

• 與程式庫程式碼類似，它可用於編寫終端使用者的程式，以搭配和不搭配發布策略運作，而無需條件邏輯。以下是說明此情況的範例程式碼片段

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.))

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 中，而 tf.keras 是 TensorFlow 針對 Keras API 規格的實作。tf.keras 是一種用於建構和訓練模型的高階 API。透過整合到 tf.keras 後端，您可以無縫地發布在 Keras 訓練架構中編寫的訓練，使用 Model.fit。

以下是您需要在程式碼中變更的內容

1. 建立適當 tf.distribute.Strategy 的執行個體。
2. 將 Keras 模型、最佳化工具和指標的建立移至 strategy.scope 內。因此，模型 call()、train_step() 和 test_step() 方法中的程式碼都將被發布，並在加速器上執行。

TensorFlow 發布策略支援所有類型的 Keras 模型—Sequential、Functional 和 子類別化

以下是用於對具有一個 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')

此範例使用 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 來取得複本數。

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 端對端整合

1. 教學課程：使用 Model.fit 和 MirroredStrategy 進行訓練。
2. 教學課程：使用 Model.fit 和 MultiWorkerMirroredStrategy 進行訓練。
3. 指南：包含使用 Model.fit 和 TPUStrategy 的範例。
4. 教學課程：使用 Model.fit 和 ParameterServerStrategy 進行參數伺服器訓練。
5. 教學課程：使用 Model.fit 和 TPUStrategy 為 GLUE 基準測試中的許多工作微調 BERT。
6. TensorFlow 模型花園 存放區，其中包含使用各種策略實作的最新模型集合。

將 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)

以上程式碼中還有一些其他需要注意的事項

1. 您使用了 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=。在沒有短批次的情況下，預設值等同於 tf.nn.compute_average_loss(..., global_batch_size=global_batch_size)，其中 global_batch_size 如上定義。（有關短批次以及如何避免或處理它們的更多資訊，請參閱自訂訓練教學課程。）

2. 您使用了 tf.nn.scale_regularization_loss 來縮放向 Model 物件註冊的正規化損失（如果有的話），縮放比例為 1/num_replicas_in_sync。對於那些與輸入相關的正規化損失，如何在每個複本 (!) 批次大小上執行平均取決於建模程式碼，而不是自訂訓練迴圈；這樣建模程式碼可以保持與複寫無關，而訓練迴圈可以保持與正規化損失的計算方式無關。

3. 當您在分發策略範圍內呼叫 apply_gradients 時，其行為會被修改。具體來說，在同步訓練期間在每個平行實例上套用梯度之前，它會執行所有複本的梯度總和。

4. 您還使用了 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))

在上面的範例中，您迭代了 dist_dataset 以提供輸入給您的訓練。您還可以使用 tf.distribute.Strategy.make_experimental_numpy_dataset 來支援 NumPy 輸入。您可以使用此 API 在呼叫 tf.distribute.Strategy.experimental_distribute_dataset 之前建立資料集。

迭代資料的另一種方法是明確地使用迭代器。當您想要針對給定的步驟數執行，而不是迭代整個資料集時，您可能會想要這樣做。現在可以修改上述迭代，首先建立一個迭代器，然後明確地在其上呼叫 next 以取得輸入資料。

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

這涵蓋了使用 tf.distribute.Strategy API 來分發自訂訓練迴圈的最簡單情況。

現在支援哪些功能？

訓練 API	MirroredStrategy	TPUStrategy	MultiWorkerMirroredStrategy	ParameterServerStrategy	CentralStorageStrategy
自訂訓練迴圈	支援	支援	支援	實驗性支援	實驗性支援

範例和教學課程

以下是一些將分發策略與自訂訓練迴圈結合使用的範例

1. 教學課程：使用自訂訓練迴圈和 MirroredStrategy 進行訓練。
2. 教學課程：使用自訂訓練迴圈和 MultiWorkerMirroredStrategy 進行訓練。
3. 指南：包含使用 TPUStrategy 的自訂訓練迴圈範例。
4. 教學課程：使用自訂訓練迴圈和 ParameterServerStrategy 進行參數伺服器訓練。
5. TensorFlow 模型花園 存放區，其中包含使用各種策略實作的最新模型集合。

其他主題

本節涵蓋了一些與多種使用案例相關的主題。

設定 TF_CONFIG 環境變數

對於多 worker 訓練，如前所述，您需要為叢集中執行的每個二進位檔設定 'TF_CONFIG' 環境變數。'TF_CONFIG' 環境變數是一個 JSON 字串，用於指定哪些任務構成叢集、它們的位址以及每個任務在叢集中的角色。tensorflow/ecosystem 儲存庫提供了一個 Kubernetes 範本，用於為您的訓練任務設定 'TF_CONFIG'。

'TF_CONFIG' 有兩個組成部分：叢集和任務。

• 叢集提供有關訓練叢集的資訊，這是一個字典，其中包含不同類型的工作，例如 worker。在多 worker 訓練中，通常有一個 worker 會承擔更多責任，例如儲存檢查點和為 TensorBoard 寫入摘要檔，以及常規 worker 所做的事情。這樣的 worker 被稱為「主導 (chief)」worker，並且習慣上將索引為 0 的 worker 指定為主導 worker（實際上，這就是 tf.distribute.Strategy 的實作方式）。
• 另一方面，任務提供有關當前任務的資訊。第一個組成部分叢集對於所有 worker 都是相同的，而第二個組成部分任務在每個 worker 上都不同，並指定該 worker 的類型和索引。

'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' 指定在 "cluster" 中有三個 worker 和兩個 "ps" 任務，以及它們的主機和埠。"task" 部分指定了當前任務在 "cluster" 中的角色—worker 1（第二個 worker）。叢集中的有效角色為 "chief"、"worker"、"ps" 和 "evaluator"。除了使用 tf.distribute.experimental.ParameterServerStrategy 之外，不應有 "ps" 工作。

下一步是什麼？

tf.distribute.Strategy 正在積極開發中。請試用它並使用 GitHub issues 提供您的意見反應。