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

用法

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

參數

fn 要運行的函數。輸出必須是 Tensor 的 tf.nest。
args (可選)fn 的位置參數。
kwargs (可選)fn 的關鍵字參數。
options (可選)tf.distribute.RunOptions 的實例，指定運行 fn 的選項。

跨副本合並fn 的返回值。返回值的結構與 fn 的返回值相同。結構中的每個元素都可以是 "per-replica" Tensor 對象或 Tensor s(例如，如果在單個副本上運行)。

使用給定的參數在每個副本上運行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)

相關用法

注：本文由純淨天空篩選整理自tensorflow.org大神的英文原創作品 tf.compat.v1.distribute.experimental.TPUStrategy.run。非經特殊聲明，原始代碼版權歸原作者所有，本譯文未經允許或授權，請勿轉載或複製。

用法

參數

返回