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

TPU 和 TPU Pod 的同步訓練。

繼承自：Strategy

用法

tf.distribute.experimental.TPUStrategy(
    tpu_cluster_resolver=None, device_assignment=None
)

要構造 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.experimental.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 內調用。