Python tf.distribute.TPUStrategy用法及代碼示例

TPU 和 TPU Pod 的同步訓練。

繼承自：Strategy

用法

tf.distribute.TPUStrategy(
    tpu_cluster_resolver=None, experimental_device_assignment=None,
    experimental_spmd_xla_partitioning=False
)

要構造 TPUStrategy 對象，您需要運行如下初始化代碼：

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)

在使用分發策略時，在策略範圍內創建的變量將在所有副本中複製，並且可以使用 all-reduce 算法保持同步。

要在 TPU 上運行 TF2 程序，您可以將 tf.keras 中的 .compile 和 .fit API 與 TPUStrategy 一起使用，或者通過直接調用 strategy.run 來編寫自己的自定義訓練循環。請注意，TPUStrategy 不支持純粹的 Eager 執行，因此請確保傳遞給 strategy.run 的函數是 tf.function 或 strategy.run 如果啟用了 Eager 行為，則在 tf.function 內調用。在 https://www.tensorflow.org/guide/tpu 中查看更多詳細信息。

在編寫自己的訓練循環時，distribute_datasets_from_function 和 experimental_distribute_dataset API 可用於在 TPU 工作人員之間分配數據集。如果您使用 tf.keras.Model 中可用的 fit 和 compile 方法，那麽 Keras 將為您處理分發。

在 TPU 上編寫自定義訓練循環的示例：

with strategy.scope():
  model = tf.keras.Sequential([
    tf.keras.layers.Dense(2, input_shape=(5,)),
  ])
  optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)

def dataset_fn(ctx):
  x = np.random.random((2, 5)).astype(np.float32)
  y = np.random.randint(2, size=(2, 1))
  dataset = tf.data.Dataset.from_tensor_slices((x, y))
  return dataset.repeat().batch(1, drop_remainder=True)
dist_dataset = strategy.distribute_datasets_from_function(
    dataset_fn)
iterator = iter(dist_dataset)

@tf.function()
def train_step(iterator):

  def step_fn(inputs):
    features, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(features, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits)

    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

  strategy.run(step_fn, args=(next(iterator),))

train_step(iterator)

對於模型並行等高級用例，您可以在創建 TPUStrategy 時設置 experimental_device_assignment 參數以指定副本數和邏輯設備數。下麵是一個使用 2 個邏輯設備和 1 個副本初始化 TPU 係統的示例。

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 1, 1, 2],
    num_replicas=1)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)

然後您可以僅在邏輯設備 0 上運行 tf.add 操作。

@tf.function()
def step_fn(inputs):
  features, _ = inputs
  output = tf.add(features, features)

  # Add operation will be executed on logical device 0.
  output = strategy.experimental_assign_to_logical_device(output, 0)
  return output
dist_dataset = strategy.distribute_datasets_from_function(
    dataset_fn)
iterator = iter(dist_dataset)
strategy.run(step_fn, args=(next(iterator),))

experimental_spmd_xla_partitioning 啟用實驗性 XLA SPMD 函數以實現模型並行性。該標誌可以減少編譯時間和 HBM 要求。在此模式下運行時，每個輸入張量必須被分區(通過 strategy.experimental_split_to_logical_devices )或完全複製(通過 strategy.experimental_replicate_to_logical_devices )到所有邏輯設備。在此模式下調用 strategy.experimental_assign_to_logical_device 將導致 ValueError。