 在 TensorFlow.org 上檢視
|
 在 Google Colab 中執行
|
 在 GitHub 上檢視原始碼
|
 下載筆記本
|
總覽
本教學課程示範如何將 tf.distribute.Strategy 用於搭配 tf.estimator 進行分散式多工作站訓練。如果您使用 tf.estimator 撰寫程式碼，而且您有興趣擴展到高效能的單一機器以外的規模，本教學課程非常適合您。
開始之前，請先閱讀分散策略指南。多 GPU 訓練教學課程也相關，因為本教學課程使用相同的模型。
設定
首先，設定 TensorFlow 和必要的匯入項目。
import tensorflow_datasets as tfds
import tensorflow as tf
import os, json
tf.compat.v1.disable_eager_execution()
輸入函式
本教學課程使用來自 TensorFlow Datasets 的 MNIST 資料集。此處的程式碼與多 GPU 訓練教學課程類似，但有一個主要差異：使用 Estimator 進行多工作站訓練時，必須依工作站數量分割資料集，以確保模型收斂。輸入資料會依工作站索引分割，讓每個工作站處理資料集的 1/num_workers 個不同部分。
BUFFER_SIZE = 10000
BATCH_SIZE = 64
def input_fn(mode, input_context=None):
datasets, info = tfds.load(name='mnist',
with_info=True,
as_supervised=True)
mnist_dataset = (datasets['train'] if mode == tf.estimator.ModeKeys.TRAIN else
datasets['test'])
def scale(image, label):
image = tf.cast(image, tf.float32)
image /= 255
return image, label
if input_context:
mnist_dataset = mnist_dataset.shard(input_context.num_input_pipelines,
input_context.input_pipeline_id)
return mnist_dataset.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
達成收斂的另一個合理方法是在每個工作站使用不同的種子隨機排序資料集。
多工作站設定
本教學課程 (相較於多 GPU 訓練教學課程) 的主要差異之一是多工作站設定。TF_CONFIG 環境變數是指定叢集組態給叢集中每個工作站的標準方式。
TF_CONFIG 有兩個元件：cluster 和 task。cluster 提供關於整個叢集的資訊，也就是叢集中的工作站和參數伺服器。task 提供關於目前工作任務的資訊。第一個元件 cluster 對於叢集中的所有工作站和參數伺服器都相同，第二個元件 task 在每個工作站和參數伺服器上都不同，並指定其自己的 type 和 index。在本範例中，工作任務 type 是 worker，工作任務 index 是 0。
為了說明目的，本教學課程示範如何在 localhost 上使用 2 個工作站設定 TF_CONFIG。實際上，您會在外部 IP 位址和連接埠上建立多個工作站，並在每個工作站上適當設定 TF_CONFIG，也就是修改工作任務 index。
os.environ['TF_CONFIG'] = json.dumps({
'cluster': {
'worker': ["localhost:12345", "localhost:23456"]
},
'task': {'type': 'worker', 'index': 0}
})
定義模型
撰寫訓練的層、最佳化工具和損失函式。本教學課程使用 Keras 層定義模型，類似於多 GPU 訓練教學課程。
LEARNING_RATE = 1e-4
def model_fn(features, labels, mode):
model = tf.keras.Sequential([
tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10)
])
logits = model(features, training=False)
if mode == tf.estimator.ModeKeys.PREDICT:
predictions = {'logits': logits}
return tf.estimator.EstimatorSpec(labels=labels, predictions=predictions)
optimizer = tf.compat.v1.train.GradientDescentOptimizer(
learning_rate=LEARNING_RATE)
loss = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(labels, logits)
loss = tf.reduce_sum(loss) * (1. / BATCH_SIZE)
if mode == tf.estimator.ModeKeys.EVAL:
return tf.estimator.EstimatorSpec(mode, loss=loss)
return tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
train_op=optimizer.minimize(
loss, tf.compat.v1.train.get_or_create_global_step()))
MultiWorkerMirroredStrategy
若要訓練模型，請使用 tf.distribute.experimental.MultiWorkerMirroredStrategy 的執行個體。MultiWorkerMirroredStrategy 會在所有工作站的所有裝置上建立模型層中所有變數的副本。它使用 CollectiveOps (一種用於集體通訊的 TensorFlow 運算子) 來彙總梯度，並讓變數保持同步。tf.distribute.Strategy 指南中有關於此策略的更多詳細資訊。
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
訓練和評估模型
接下來，在 Estimator 的 RunConfig 中指定分散策略，然後透過叫用 tf.estimator.train_and_evaluate 進行訓練和評估。本教學課程僅透過 train_distribute 指定策略來分散訓練。也可以透過 eval_distribute 分散評估。
config = tf.estimator.RunConfig(train_distribute=strategy)
classifier = tf.estimator.Estimator(
model_fn=model_fn, model_dir='/tmp/multiworker', config=config)
tf.estimator.train_and_evaluate(
classifier,
train_spec=tf.estimator.TrainSpec(input_fn=input_fn),
eval_spec=tf.estimator.EvalSpec(input_fn=input_fn)
)
最佳化訓練效能
您現在有一個模型和一個由 tf.distribute.Strategy 強化的多工作站功能 Estimator。您可以嘗試下列技術來最佳化多工作站訓練的效能
- 增加批次大小： 此處指定的批次大小是每個 GPU。一般而言，建議使用可容納 GPU 記憶體的最大批次大小。
- 轉換變數： 如果可能，將變數轉換為
tf.float。官方 ResNet 模型包含如何執行此操作的範例。 使用集體通訊：
MultiWorkerMirroredStrategy提供多種集體通訊實作。RING使用 gRPC 作為跨主機通訊層，實作以環狀為基礎的集體運算。NCCL使用 Nvidia 的 NCCL 來實作集體運算。AUTO將選擇權延後至執行階段。
集體實作的最佳選擇取決於 GPU 的數量和種類，以及叢集中的網路互連。若要覆寫自動選擇，請將有效值指定給
MultiWorkerMirroredStrategy建構函式的communication參數，例如communication=tf.distribute.experimental.CollectiveCommunication.NCCL。
請造訪指南中的效能章節，以進一步瞭解其他策略，以及可用於最佳化 TensorFlow 模型效能的工具。