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

Question

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

以下是获取Loss的代码：

    # 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 是数组，如何将这些数组合并为一个数组。

非常感谢。

标签： tensorflowgpudistributiontpu

Answer 1

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