Python tf.distribute.TPUStrategy.run用法及代码示例

用法

run(
    fn, args=(), kwargs=None, options=None
)

在每个 TPU 副本上运行由 fn 定义的计算。

在每个副本上执行 fn 指定的操作。如果 args 或 kwargs 具有 tf.distribute.DistributedValues ，例如由 tf.distribute.Strategy.experimental_distribute_dataset 或 tf.distribute.Strategy.distribute_datasets_from_function 中的 tf.distribute.DistributedDataset 生成的那些，当 fn 在特定副本上执行时，它将使用以下组件执行tf.distribute.DistributedValues 对应于该副本。

fn 可以调用 tf.distribute.get_replica_context() 来访问成员，例如 all_reduce 。

args 或 kwargs 中的所有参数应该是张量嵌套或包含张量或复合张量的 tf.distribute.DistributedValues。

示例用法：

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)
@tf.function
def run():
  def value_fn(value_context):
    return value_context.num_replicas_in_sync
  distributed_values = (
      strategy.experimental_distribute_values_from_function(value_fn))
  def replica_fn(input):
    return input * 2
  return strategy.run(replica_fn, args=(distributed_values,))
result = run()