Python tf.distribute.experimental.ValueContext用法及代碼示例

一個類包裝分發函數所需的信息。

用法

tf.distribute.experimental.ValueContext(
    replica_id_in_sync_group=0, num_replicas_in_sync=1
)

這是一個上下文類，它傳遞給strategy.experimental_distribute_values_from_function 中的value_fn，並包含有關計算副本的信息。 num_replicas_in_sync 和replica_id 可用於自定義每個副本的值。

示例用法：

1. 直接構建。

def value_fn(context):
  return context.replica_id_in_sync_group/context.num_replicas_in_sync
context = tf.distribute.experimental.ValueContext(
  replica_id_in_sync_group=2, num_replicas_in_sync=4)
per_replica_value = value_fn(context)
per_replica_value
0.5

1. 由 experimental_distribute_values_from_function 傳入。

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
def value_fn(value_context):
  return value_context.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)