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

用法

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

使用给定的参数在每个副本上运行fn。

在每个副本上执行 fn 指定的操作。如果args 或kwargs 具有"per-replica" 值，例如由“分布式Dataset”产生的值，则在特定副本上执行fn 时，它将与那些"per-replica" 的组件一起执行对应于该副本的值。

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

args 或 kwargs 中的所有参数应该是张量嵌套或包含张量或复合张量的 per-replica 对象。

用户可以将策略特定选项传递给options 参数。在TPUStrategy.run 中启用分桶动态形状的示例是：

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)

options = tf.distribute.RunOptions(
    experimental_bucketizing_dynamic_shape=True)

dataset = tf.data.Dataset.range(
   strategy.num_replicas_in_sync, output_type=dtypes.float32).batch(
       strategy.num_replicas_in_sync, drop_remainder=True)
input_iterator = iter(strategy.experimental_distribute_dataset(dataset))

@tf.function()
def step_fn(inputs):
 output = tf.reduce_sum(inputs)
 return output

strategy.run(step_fn, args=(next(input_iterator),), options=options)