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

用法

experimental_distribute_values_from_function(
    value_fn
)

从 value_fn 生成 tf.distribute.DistributedValues 。

此函数用于生成 tf.distribute.DistributedValues 以传递给 run , reduce 或其他在不使用数据集时采用分布式值的方法。

示例用法：

1. 每个副本返回常量值：

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
def value_fn(ctx):
  return tf.constant(1.)
distributed_values = (
     strategy.experimental_distribute_values_from_function(
       value_fn))
local_result = strategy.experimental_local_results(distributed_values)
local_result
(<tf.Tensor:shape=(), dtype=float32, numpy=1.0>,
 <tf.Tensor:shape=(), dtype=float32, numpy=1.0>)

1. 根据replica_id在数组中分配值：

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
array_value = np.array([3., 2., 1.])
def value_fn(ctx):
  return array_value[ctx.replica_id_in_sync_group]
distributed_values = (
     strategy.experimental_distribute_values_from_function(
       value_fn))
local_result = strategy.experimental_local_results(distributed_values)
local_result
(3.0, 2.0)

1. 使用 num_replicas_in_sync 指定值：

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
def value_fn(ctx):
  return ctx.num_replicas_in_sync
distributed_values = (
     strategy.experimental_distribute_values_from_function(
       value_fn))
local_result = strategy.experimental_local_results(distributed_values)
local_result
(2, 2)

1. 在设备上放置值并分发：

strategy = tf.distribute.TPUStrategy()
worker_devices = strategy.extended.worker_devices
multiple_values = []
for i in range(strategy.num_replicas_in_sync):
  with tf.device(worker_devices[i]):
    multiple_values.append(tf.constant(1.0))

def value_fn(ctx):
  return multiple_values[ctx.replica_id_in_sync_group]

distributed_values = strategy.
  experimental_distribute_values_from_function(
  value_fn)