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

用法

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

参数

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

跨副本合并fn 的返回值。返回值的结构与 fn 的返回值相同。结构中的每个元素都可以是 tf.distribute.DistributedValues 、 Tensor 对象或 Tensor s(例如，如果在单个副本上运行)。

在每个 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()

相关用法

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

用法

参数

返回

示例用法：