Logging in `on_test_epoch_end` with multiple dataloaders

[pschroeppel]

Bug description

Hello,
I am running test() with multiple dataloaders and do the logging in on_test_epoch_end(). I basically follow #16520 regarding the "Example with two DataLoaders".

In my case, I log different metrics for the different dataloaders (because each dataloader corresponds to a set set and I directly put the test set name into the logged name). For example:

• dataloader 1 logs clevr3d_cropped/ari
• dataloader 2 logs clevr3d/ari

However, the final table that is printed by PyTorch-Lightning, shows each logged value duplicated for each of the two dataloaders, which is incorrect:

        Test metric               DataLoader 0               DataLoader 1
─────────────────────────────
     clevr3d/ari         0.6213954091072083         0.6213954091072083
 clevr3d_cropped/ari      0.5558694005012512         0.5558694005012512
─────────────────────────────

Correctly, the table should look as follows:

        Test metric               DataLoader 0               DataLoader 1
─────────────────────────────
     clevr3d/ari         0.6213954091072083                -
 clevr3d_cropped/ari           -                          0.5558694005012512
─────────────────────────────

I think the problem is that self.log() does not know for which dataloader the log is and therefore logs for both dataloaders.

I also tried appending the dataloader index in the log, as described in the example mentioned above, but this didn't change anything.

My question now is: is this expected behaviour, or is it a bug, or am I using PyTorch-Lightning incorrectly?

Thanks and best,
Philipp

What version are you seeing the problem on?

v2.5

Reproduced in studio

No response

How to reproduce the bug

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):

More info

No response

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions - the Lightning Team!

[matsumotosan]

@pschroeppel I tested out a simple example and found the logging feature to output the expected table. Could you clarify if there is something different that you are trying to do?

I tried both self.log and self.log_dict and found both to work.

import tempfile
from typing import Any, Optional

import torch
from torch import Tensor
from torch.utils.data import DataLoader, Dataset

import lightning as L
from lightning import Trainer
from lightning.pytorch import seed_everything


def main():
    with tempfile.TemporaryDirectory() as tempdir:
        seed_everything(42, workers=True)
        model = BoringModel()
        trainer = Trainer(default_root_dir=tempdir, max_epochs=5, deterministic=True)

        dm = BoringDataModule()
        trainer.fit(model, dm)
        trainer.test(model, dm)


class RandomDataset(Dataset):
    def __init__(self, size: int, length: int):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index: int) -> Tensor:
        return self.data[index]

    def __len__(self) -> int:
        return self.len


class BoringDataModule(L.LightningDataModule):
    def __init__(self) -> None:
        super().__init__()

    def setup(self, stage: str) -> None:
        if stage == "fit":
            self.random_train = RandomDataset(32, 8)
        if stage in ("fit", "validate"):
            self.random_val = RandomDataset(32, 8)
        if stage == "test":
            self.random_test1 = RandomDataset(32, 8)
            self.random_test2 = RandomDataset(32, 8)

    def train_dataloader(self) -> DataLoader:
        return DataLoader(self.random_train)

    def val_dataloader(self) -> DataLoader:
        return DataLoader(self.random_val)

    def test_dataloader(self):
        return [DataLoader(self.random_test1), DataLoader(self.random_test2)]


class BoringModel(L.LightningModule):
    def __init__(self) -> None:
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x: Tensor) -> Tensor:
        return self.layer(x)

    def loss(self, preds: Tensor, labels: Optional[Tensor] = None) -> Tensor:
        if labels is None:
            labels = torch.ones_like(preds)
        return torch.nn.functional.mse_loss(preds, labels)

    def step(self, batch: Any) -> Tensor:
        output = self(batch)
        return self.loss(output)

    def training_step(self, batch: Any, batch_idx: int):
        return {"loss": self.step(batch)}

    def validation_step(self, batch: Any, batch_idx: int):
        return {"loss": self.step(batch)}

    def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0):
        loss = self.step(batch)
        if dataloader_idx == 0:
            self.log("test_loss0", loss)
            # self.log_dict({"test_loss0": loss})
        else:
            self.log("test_loss1", loss)
            # self.log_dict({"test_loss1": loss})

    def configure_optimizers(self):
        return torch.optim.SGD(self.parameters(), lr=0.1)


if __name__ == "__main__":
    main()

The trainer.test() returns the table below:

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃        Test metric        ┃       DataLoader 0        ┃       DataLoader 1        ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│        test_loss0         │      320.73876953125      │                           │
│        test_loss1         │                           │    255.02682495117188     │
└───────────────────────────┴───────────────────────────┴───────────────────────────┘