A little composition wrapper for Pytorch loss functions. It calculates weighted sum of multiple loss functions.

weighted_sum_composition_loss.py

from collections.abc import Collection

from torch import nn


class WeightedSumCompositionLoss(nn.Module):
    def __init__(self, loss_funcs: Collection, weights: Collection):
        super().__init__()
        self.loss_funcs = loss_funcs
        self.weights = weights

    def forward(self, output, target):
        loss = 0.
        for idx in range(len(self.loss_funcs)):
            loss += self.loss_funcs[idx].forward(output, target) * self.weights[idx]
        return loss
 

A little composition wrapper for Pytorch loss functions. It calculates weighted sum of multiple loss functions.

Example how to use:

loss_fn = WeightedSumCompositionLoss(
            loss_funcs=[
                VGGPerceptualLoss(resize=True),
                torch.nn.BCEWithLogitsLoss(reduction='mean'),
            ],
            weights=[1, 1],
        )