Logging for elements of compound/weighted loss

[maciejzj]

Hi, I have a pretty standard LightningModule where the loss is passed to the Module via init, like here:

class GenericImgToImgTask(pl.LightningModule):
    def __init__(
        self,
        model: torch.nn.Module,
        loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
    ):
        """Generic task (Lightning module) for image-to-image scenarios.

        :param model: PyTorch model (Module) to wrap around.
        :param loss: Training loss (defaults to `torch.nn.L1Loss`).
        """
        super().__init__()
        self._model = model
        self._loss = loss if loss is not None else L1Loss()

The implementation of train/val/test steps is pretty standard for a supervised training scheme, e.g.:

@override
def training_step(
    self, train_batch: torch.Tensor, batch_idx: int
) -> torch.Tensor | dict[str, torch.Tensor]:
    x, y_gt = train_batch
    y_pred = self._model(x)
    loss = self._loss(y_pred, y_gt)
    self.log('train_loss', loss, on_step=False, on_epoch=True, prog_bar=True)
    return loss

Now, I'd like to be able to inject a more complex loss with multiple weighted components, something like this:

class WeightedLoss(torch.nn.Module):
    def __init__(
        self,
        losses: Sequence[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]],
        weights: Sequence[float],
    ):
        super().__init__()
        self._losses = torch.nn.ModuleList(losses)
        self._weights = weights

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        score = 0.0
        for loss, weight in zip(self._losses, self._weights):
            score += loss(pred, target) * weight
        return score / sum(self._weights)

So finally I can do something like:

GenericImgToImgTask(MyPyTorchModel(), WeightedLoss([L1Loss()], [BCELoss()]), [0.25, 0.75])

One problem is that I want to log/monitor each of the losses in the weighted loss. I have figured out that I can pass a pl.loggers.Logger to the weighted loss and call self._logger.log_metrics with results of the loss function inside the for loop. However, this isn't as great as self.log inside a LightningModule (I cannot use on_epoch and logging step is rather problematic). I thought about making WeightedLoss a LightningModule, but then I would have to connect it to the Lightning's experiment somehow. Is there an elegant way to solve this problem?