I am working on a framework on top of tensorflow that can train models in a generic way (there's more to it but that's not relevant for the question).
User will need to specify his compute_loss function which should return a dictionary loss name -> loss value. Even though models are optimised on one loss, this is to keep track of eventual intermediary losses, essentially for logging.
def compute_loss(**data_dict, **predictions):
# compute some losses
return { 'total_loss': total_loss, 'intermediary_loss1': intermediary_loss1 }
The total_loss key has to exists and it references the loss used for optimisation.
Then I define a generic train_step:
def train_step(data_dict):
with tf.GradientTape() as tape:
predictions = model(**data_dict, training=True)
loss_dict = compute_loss(**data_dict, **predictions)
total_loss = loss_dict['total_loss']
gradients = tape.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return loss_dict
However when using using MirroredStrategy for multi-gpu training, I need to wrap the train_step as so (from TF doc):
@tf.function
def distributed_train_step(dist_inputs):
per_replica_losses = mirrored_strategy.run(train_step, args=(dist_inputs,))
return mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
axis=None)
I need to apply a reduce to aggregate losses over replicas, but how can I apply a reduce operation if my train step returns a dict?
The underlying question behind this is how and where in the code can I keep track of losses name and values at the same time after execution of a train step.
Thanks