当前位置: 首页>>代码示例 >>用法及示例精选 >>正文


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

用法

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

参数

  • fn 要运行的函数。输出必须是 Tensortf.nest
  • args (可选)fn 的位置参数。
  • kwargs (可选)fn 的关键字参数。
  • options (可选)tf.distribute.RunOptions 的实例,指定运行 fn 的选项。

返回

  • 跨副本合并fn 的返回值。返回值的结构与 fn 的返回值相同。结构中的每个元素都可以是 "per-replica" Tensor 对象或 Tensor s(例如,如果在单个副本上运行)。

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

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

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

argskwargs 中的所有参数应该是张量嵌套或包含张量或复合张量的 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)

相关用法


注:本文由纯净天空筛选整理自tensorflow.org大神的英文原创作品 tf.compat.v1.distribute.experimental.TPUStrategy.run。非经特殊声明,原始代码版权归原作者所有,本译文未经允许或授权,请勿转载或复制。