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

用法

reduce(
    reduce_op, value, axis
)

跨副本减少 value。

给定 run 返回的 per-replica 值，比如 per-example 损失，批次将被划分到所有副本中。此函数允许您跨副本进行聚合，也可以选择跨批次元素进行聚合。例如，如果全局批量大小为 8 个副本和 2 个副本，则示例 [0, 1, 2, 3] 的值将位于副本 0 上，而 [4, 5, 6, 7] 将位于副本 1 上。默认情况下，reduce 将仅跨副本聚合，返回[0+4, 1+5, 2+6, 3+7] 。当每个副本计算标量或其他没有"batch" 维度的值(如梯度)时，这很有用。更常见的情况是，您希望在全局批次中进行聚合，您可以通过将批次维度指定为 axis 来获得，通常是 axis=0 。在这种情况下，它将返回一个标量 0+1+2+3+4+5+6+7 。

如果有最后一个部分批次，您将需要指定一个轴，以便生成的形状在副本之间是一致的。因此，如果最后一批的大小为 6，并且它被分为 [0, 1, 2, 3] 和 [4, 5]，除非您指定 axis=0 ，否则您将得到形状不匹配。如果您指定 tf.distribute.ReduceOp.MEAN ，则使用 axis=0 将使用正确的分母 6。将此与计算 reduce_mean 以获取每个副本的标量值进行对比，此函数用于平均这些平均值，这将权衡一些值 1/8和其他 1/4 。

例如：

strategy = tf.distribute.experimental.CentralStorageStrategy(
    compute_devices=['CPU:0', 'GPU:0'], parameter_device='CPU:0')
ds = tf.data.Dataset.range(10)
# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(ds)

with strategy.scope():
  @tf.function
  def train_step(val):
    # pass through
    return val

  # Iterate over the distributed dataset
  for x in dist_dataset:
    result = strategy.run(train_step, args=(x,))

result = strategy.reduce(tf.distribute.ReduceOp.SUM, result,
                         axis=None).numpy()
# result:array([ 4,  6,  8, 10])

result = strategy.reduce(tf.distribute.ReduceOp.SUM, result, axis=0).numpy()
# result:28