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 ?
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.
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.