2

I'm trying to use strategy in Tensorflow. I know how I can get summation and mean value of the result from each replica. However, if the result of each replica is array, how can I merge these arrays from each replica to one array?

The following is the code for getting 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

I used the function strategy.reduce() to get the summation of loss.

If the per_replica_loss is array, how can I merge these arrays to one array.

Thanks a lot.

MrJasonLi
  • 21
  • 3

1 Answers1

0

Use strategy.unwrap(tensor) . It return a tuple of each replica PerReplica

Jude TCHAYE
  • 434
  • 5
  • 14