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 内调用。