Python tf.compat.v1.distribute.MirroredStrategy.reduce用法及代碼示例

用法

reduce(
    reduce_op, value, axis=None
)

跨副本減少 value 並在當前設備上返回結果。

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
def step_fn():
  i = tf.distribute.get_replica_context().replica_id_in_sync_group
  return tf.identity(i)

per_replica_result = strategy.run(step_fn)
total = strategy.reduce("SUM", per_replica_result, axis=None)
total
<tf.Tensor:shape=(), dtype=int32, numpy=1>

要查看多個副本的情況，請考慮具有 2 個 GPU 的 MirroredStrategy 的相同示例：

strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
def step_fn():
  i = tf.distribute.get_replica_context().replica_id_in_sync_group
  return tf.identity(i)

per_replica_result = strategy.run(step_fn)
# Check devices on which per replica result is:
strategy.experimental_local_results(per_replica_result)[0].device
# /job:localhost/replica:0/task:0/device:GPU:0
strategy.experimental_local_results(per_replica_result)[1].device
# /job:localhost/replica:0/task:0/device:GPU:1

total = strategy.reduce("SUM", per_replica_result, axis=None)
# Check device on which reduced result is:
total.device
# /job:localhost/replica:0/task:0/device:CPU:0

此 API 通常用於匯總從不同副本返回的結果，用於報告等。例如，可以在打印之前使用此 API 對從不同副本計算的損失進行平均。

注意：結果被複製到"current" 設備——它通常是運行程序的工作線程的 CPU。對於 TPUStrategy ，它是第一個 TPU 主機。對於多客戶端 MultiWorkerMirroredStrategy ，這是每個工作人員的 CPU。

有許多不同的 tf.distribute API 用於減少副本之間的值：

• tf.distribute.ReplicaContext.all_reduce ：這與 Strategy.reduce 的不同之處在於它用於副本上下文並且不會將結果複製到主機設備。 all_reduce 通常應用於訓練步驟中的減少，例如梯度。
• tf.distribute.StrategyExtended.reduce_to 和 tf.distribute.StrategyExtended.batch_reduce_to ：這些 API 是 Strategy.reduce 的更高級版本，因為它們允許自定義結果的目標。它們也在跨副本上下文中被調用。

軸應該是什麽？

給定 run 返回的 per-replica 值，比如 per-example 損失，批次將被劃分到所有副本中。此函數允許您通過相應地指定軸參數來跨副本聚合，也可以選擇跨批處理元素聚合。

例如，如果全局批量大小為 8 和 2 個副本，則示例的值 [0, 1, 2, 3] 將位於副本 0 上，[4, 5, 6, 7] 將位於副本 1 上。axis=None , reduce 將僅跨副本聚合，返回 [0+4, 1+5, 2+6, 3+7] .當每個副本計算標量或其他沒有"batch" 維度的值(如梯度或損失)時，這很有用。

strategy.reduce("sum", per_replica_result, axis=None)

有時，您會希望在全局批次和所有副本之間進行聚合。您可以通過將批處理維度指定為 axis 來獲得此行為，通常是 axis=0 。在這種情況下，它將返回一個標量 0+1+2+3+4+5+6+7 。

strategy.reduce("sum", per_replica_result, axis=0)

如果有最後一個部分批次，您將需要指定一個軸，以便生成的形狀在副本之間是一致的。因此，如果最後一批的大小為 6，並且它被分為 [0, 1, 2, 3] 和 [4, 5]，除非您指定 axis=0 ，否則您將得到形狀不匹配。如果您指定 tf.distribute.ReduceOp.MEAN ，則使用 axis=0 將使用正確的分母 6。將此與計算 reduce_mean 以獲取每個副本的標量值進行對比，此函數用於平均這些平均值，這將權衡一些值 1/8和其他 1/4 。