Multiple Loss Tracking on Train and Evaluate Steps

Hello
I am using transformers 4.45.2
I created a model with two different tasks (so two different losses) and can train it fine. But I wanted to explicitly track both of my losses, not only single loss, during training and eval steps.
I have found some work arounds to partially solve the issue, but they don’t seem like clean solutions to me and all of them failed at some detail.
My model already returns loss_total, loss_an and loss_b and some other things in its output.
Obviously, without modification, trainer only prints out loss_total.
I want to see all 3 losses, and also on epoch end, where the evaluation happens, I want to see these values on my evaluation dataset.
On epoch end, I want to see
train_total_loss, validation_total_loss, train_loss_a, train_loss_b, validation_loss_a, validation_loss_b.

Is there a clean solution for this
Thanks

1 Like

BTW, I tried
1- customizing the trainer class
2- using callbacks
3- using compute_metrics

1 Like