Feature request: Print out values of individual Loss components in training

[xboldocky]

Hello,

I believe it might be useful to have an option to declare loss component names in trainer class, so they get be printed along with the train_loss value during training. For example:

Loss Components:

regulation_loss = 5. * ((x == ref) ^ 2)
regulation_loss.name = 'ref_tracking'
state_lower_bound_penalty = 10.*(x > xmin)
state_upper_bound_penalty = 10.*(x < xmax)

state_lower_bound_penalty.name = 'x_min'
state_upper_bound_penalty.name = 'x_max'

Trainer Class:

 trainer = Trainer(
        problem,
        train_loader, dev_loader,
        optimizer=optimizer,
        epochs=100,
        train_metric='train_loss',
        eval_metric='dev_loss',
        display=['train_loss', 'ref_tracking', 'x_min', 'x_max', 'el_time'],
    )

Displaying this information could be very insightful/informative in scenarios where training diverges. The user could easily identify which component causes divergence. In addition, it could be also informative in applications involving (self-)adaptive penalty coefficients.

This feature addition might be nontrivial to implement in the current codebase, since the Trainer class does not currently have access to individual loss component values.