Global & Special Loggers

[jamie-mcg]

Description & Motivation

The trainer currently treats all loggers the same, meaning that if you design a custom logger for a specific purpose, it gets called every time self.log is called. This means if I have two separate CSVLogger's defined, one for normal training metrics and another for some non-standard metrics that do not fit with the 1st loggers schema, then PL will still attempt to log all metrics to both CSVLoggers.

It would be nice to have the ability to separate Global loggers (i.e. ones that should be activated with self.log) and those which are more specific to an independent purpose.

Pitch

I believe it can be done fairly easily by allowing for a separate logger-type list to be added to the Trainer class (i.e. logger_others (or some better name...), then the LoggerConnector class can simply configure the two separately.

Alternatives

I think the only alternative is to instantiate a separate logger class within the Callback? But this has all sorts of integration issues that you need to do to get it to work with the trainer.

Additional context

Very happy to open a PR with a proposed solution for this.

cc @lantiga @Borda

[SkafteNicki]

Hi @jamie-mcg, thanks for raising this issue. I would say that what you are proposing is fairly doable already in lightning today without having to make changes to the codebase. All loggers are accessible through self.loggers[0], self.loggers[1], ... inside your lightningmodule and in the same way self.loggers[0].log_metrics, ... can be called to only log values to a specific logger. So just to give you an example, in the below example I initialize 3 loggers and in my lightning module have defined a sub_log method which I can then use to only log to a subset of my loggers:

import tempfile
import torch
import torch.nn as nn
from lightning.pytorch import LightningModule, Trainer
from lightning.pytorch.loggers import CSVLogger, TensorBoardLogger


class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(5, 1)
        
    def forward(self, x):
        return self.linear(x)
    
    def training_step(self, batch, batch_idx):
        x = batch[0]
        y_hat = self(x)
        y_true = torch.randn_like(y_hat)
        loss = torch.nn.functional.mse_loss(y_hat, y_true)
        
        # Standard logging to all loggers
        self.log("loss", loss)
        
        # Access individual loggers by indexing
        self.loggers[0].log_metrics({"tensorboard_metric": 0.1}, step=batch_idx)
        self.loggers[1].log_metrics({"csv_main_metric": 0.2}, step=batch_idx) 
        self.loggers[2].log_metrics({"csv_special_metric": 0.3}, step=batch_idx)
        
        # Use custom method to log to subset
        self.sub_log({"subset_metric": 0.5}, logger_indices=[0, 2])
        return loss
    
    def sub_log(self, metrics: dict, logger_indices: list[int] = None, step: int = None):
        """Log metrics to a subset of loggers."""
        if logger_indices is None:
            logger_indices = list(range(len(self.loggers)))
        
        for idx in logger_indices:
            if 0 <= idx < len(self.loggers):
                self.loggers[idx].log_metrics(metrics, step=step)
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters())


def main():
    with tempfile.TemporaryDirectory() as tmp_dir:
        loggers = [
            TensorBoardLogger(tmp_dir, name="tensorboard"),   # Logger 0
            CSVLogger(tmp_dir, name="csv_main"),              # Logger 1  
            CSVLogger(tmp_dir, name="csv_special"),           # Logger 2
        ]
        
        model = SimpleModel()
        
        dataset = torch.utils.data.TensorDataset(torch.randn(100, 5))
        dataloader = torch.utils.data.DataLoader(dataset, batch_size=16)
        
        trainer = Trainer(
            logger=loggers,
            max_epochs=1,
            limit_train_batches=3,
            enable_progress_bar=False
        )
        
        trainer.fit(model, dataloader)


if __name__ == "__main__":
    main()

[SkafteNicki]

Closing issue, but feel free to ping and reopen if necessary.