当前位置: 首页>>代码示例 >>用法及示例精选 >>正文


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

用法

reduce(
    reduce_op, value, axis=None
)

参数

  • reduce_op 一个 tf.distribute.ReduceOp 值,指定如何组合值。允许使用枚举的字符串表示,例如"SUM"、"MEAN"。
  • value tf.distribute.DistributedValues 实例,例如由 Strategy.run 返回,组合成一个张量。当与OneDeviceStrategy 或默认策略一起使用时,它也可以是常规张量。
  • axis 指定要在每个副本的张量内减少的维度。通常应设置为批量维度,或 None 以仅跨副本减少(例如,如果张量没有批量维度)。

返回

  • 一个Tensor

跨副本减少 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 用于减少副本之间的值:

轴应该是什么?

给定 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

相关用法


注:本文由纯净天空筛选整理自tensorflow.org大神的英文原创作品 tf.compat.v1.distribute.experimental.TPUStrategy.reduce。非经特殊声明,原始代码版权归原作者所有,本译文未经允许或授权,请勿转载或复制。