how to convert perreplica to tensor?

Question

When training with multi gpu in tensorflow2.0, perreplica would be reduce by below code:

strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)

However, if I just want to collect(no 'sum reduce' or 'mean reduce') all gpu's predictions into a tensor:

per_replica_losses, per_replica_predicitions = strategy.experimental_run_v2(train_step, args=(dataset_inputs,))
# how to convert per_replica_predicitions to a tensor ?

Answer 1

In short, you can convert PerReplica result into a tuple of tensors like this:

tensors_tuple = per_replica_predicitions.values

the return tensors_tuple will be a tuple of predictions from each replicas/devices:

(predicton_tensor_from_dev0, prediction_tensor_from_dev1,...)

The number of elements in this tuple is determined by your devices available to the distributed strategy. Specially, if the strategy runs on a single replica/device, the return value from strategy.experimental_run_v2 will be the same as calling train_step function directly (tensor or list of tensors decided by your train_step). So you might want to write the code like this:

per_replica_losses, per_replica_predicitions = strategy.experimental_run_v2(train_step, args=(dataset_inputs,))

if strategy.num_replicas_in_sync > 1:
    predicition_tensors = per_replica_predicitions.values
else:
    predicition_tensors = per_replica_predicitions

PerReplica is a class object wrapping the results of distributed running. You can find its definition here, there are more properties/methods for us to operate the PerReplica object.