如何在Tensorflow 2中合并策略的结果

Question

我正在尝试在Tensorflow中使用策略。我知道如何从每个副本中获得结果的总和和平均值。但是，如果每个复制副本的结果都是阵列，我如何将这些阵列从每个复制副本合并到一个阵列？

以下是获取损失的代码：

# val
def distributed_val(ds):
total_loss = tf.cast(0.0, tf.float32)
num_train_batches = tf.cast(0.0, tf.float32)
for one_batch in ds:
per_replica_loss = strategy.experimental_run_v2(
self.loss, args=(one_batch,))
total_loss += strategy.reduce(
tf.distribute.ReduceOp.SUM, per_replica_loss, axis=None)
num_train_batches += 1
return total_loss, num_train_batches

我使用了函数strategy.reduce((来获得损失的总和。

如果per_replica_loss是数组，我如何将这些数组合并为一个数组。

非常感谢。

Answer 1

使用strategy.unwrap(张量(。它返回每个副本PerReplica 的元组