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

用法

reduce(
    reduce_op, value, axis
)

跨副本减少 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 。