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

設備列表上的狀態和計算分布策略。

用法

tf.distribute.Strategy(
    extended
)

有關概述和示例，請參閱指南。請參閱 tf.distribute.StrategyExtended 和 tf.distribute 以獲取此頁麵上提到的概念詞匯表，例如 "per-replica"、replica 和 reduce。

簡而言之：

• 要將其與 Keras compile /fit 一起使用，請閱讀。
• 您可以將 tf.distribute.Strategy 的後代傳遞給 tf.estimator.RunConfig 以指定 tf.estimator.Estimator 應如何分配其計算。見指南。
• 否則，使用tf.distribute.Strategy.scope 指定在構建執行模型時應使用的策略。 (這會將您置於此策略的“cross-replica 上下文”中，這意味著該策略可以控製變量放置等事情。)

• 如果您正在編寫自定義訓練循環，則需要調用更多方法，請參閱指南：

  • 首先正常創建tf.data.Dataset。
  • 使用 tf.distribute.Strategy.experimental_distribute_dataset 將 tf.data.Dataset 轉換為產生 "per-replica" 值的東西。如果要手動指定數據集應如何跨副本分區，請改用tf.distribute.Strategy.distribute_datasets_from_function。
  • 使用 tf.distribute.Strategy.run 對每個副本運行一次函數，獲取可能是 "per-replica" 的值(例如，來自 tf.distribute.DistributedDataset 對象)並返回 "per-replica" 值。該函數在"replica context"中執行，這意味著每個操作在每個副本上單獨執行。
  • 最後使用一種方法(例如 tf.distribute.Strategy.reduce )將生成的 "per-replica" 值轉換為普通的 Tensor 。

自定義訓練循環可以很簡單：

with my_strategy.scope():
  @tf.function
  def distribute_train_epoch(dataset):
    def replica_fn(input):
      # process input and return result
      return result

    total_result = 0
    for x in dataset:
      per_replica_result = my_strategy.run(replica_fn, args=(x,))
      total_result += my_strategy.reduce(tf.distribute.ReduceOp.SUM,
                                         per_replica_result, axis=None)
    return total_result

  dist_dataset = my_strategy.experimental_distribute_dataset(dataset)
  for _ in range(EPOCHS):
    train_result = distribute_train_epoch(dist_dataset)

這需要一個普通的 dataset 和 replica_fn 並使用上麵名為 my_strategy 的特定 tf.distribute.Strategy 運行它。 replica_fn 中創建的任何變量都是使用 my_strategy 的策略創建的，並且 replica_fn 調用的庫函數可以使用 get_replica_context() API 來實現 distributed-specific 行為。

您可以使用 reduce API 跨副本聚合結果，並將其用作 tf.distribute.DistributedDataset 上一次迭代的返回值。或者您可以使用tf.keras.metrics(例如損失、準確性等)在給定時期內跨步驟累積指標。

有關更詳細的示例，請參閱自定義訓練循環教程。

注意：tf.distribute.Strategy 目前不支持 TensorFlow 的分區變量(單個變量跨多個設備拆分)。