Support custom weighted loss in gradient accumulation
Is your feature request related to a problem? Please describe.
When using triplet loss, there could be some "easy triplets" contributing 0.0 to the batch loss.
Currently, pytorch-lighting assumes all samples in a mini-batch do contribute non-zero value to the closure_loss. This should be an issue for online triplet mining since any arbitrary portion of a batch could contribute 0.0 to the batch loss. Simply re-weight the batch loss by 1 / self.accumulate_grad_batches could lead to gradient oscillation.
https://github.com/williamFalcon/pytorch-lightning/blob/62f6f92fdf6343e17336f0a28247ff95d35bac3a/pytorch_lightning/trainer/train_loop_mixin.py#L352-L354
Describe the solution you'd like
Support returning an optional callback metric batch_loss_efficiency_ratio and re-weight the loss based on this factor. I have implemented a demo for this idea and could submit a PR if you think this is worth following.
huntzhan@9d8b78c
And the usage:
In __init__
    self.loss = nn.TripletMarginLoss(margin=self.hparams.triplet_loss_margin, reduction='none')

In training_step
    sample_losses = self.meta_step(batch, batch_idx)
    non_zero_mask = (sample_losses != 0)
    batch_loss_efficiency_ratio = float(non_zero_mask.sum()) / sample_losses.shape[0]
    if batch_loss_efficiency_ratio > 0:
        loss = sample_losses[non_zero_mask].mean()
    else:
        loss = torch.tensor(0.0, requires_grad=True)

    tqdm_dict = {
            'train_loss': loss,
            'batch_loss_efficiency_ratio': batch_loss_efficiency_ratio,
    }
    return {
            'loss': loss,
            'progress_bar': tqdm_dict,
            'log': tqdm_dict,
    }

Describe alternatives you've considered
N/A
Additional context
N/A