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