[Bug] Logging Separately Per Epoch

[justusschock]

🐛 Bug

Logs separate plots for each epoch

To Reproduce

import os
import torch
from torch.utils.data import Dataset, DataLoader
from pytorch_lightning import LightningModule, Trainer
from torch.nn import functional as F
from torchmetrics.classification import Accuracy

class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)
    def __getitem__(self, index):
        return self.data[index], 1
    def __len__(self):
        return self.len

class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)
        self.loss_fn = torch.nn.CrossEntropyLoss()
        self.train_accuracy = Accuracy()
        self.val_accuracy = Accuracy()

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_id):
        return self.step(batch, step_type="train")

    def validation_step(self, batch, batch_id):
        return self.step(batch, step_type="val")

    def step(self, batch, step_type: str):
        data, y = batch
        y_pred = self(data)
        loss = self.loss_fn(y_pred, y)
        self.log(f"{step_type}_loss", loss, prog_bar=True, sync_dist=False)
        # Compute accuracy
        y_pred = F.softmax(y_pred, dim=1)
        self.__getattr__(f"{step_type}_accuracy")(y_pred, y)
        self.log(f"{step_type}_accuracy", self.__getattr__(f"{step_type}_accuracy"), on_step=True, on_epoch=True)
        return loss

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

def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    val_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    test_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=5,
        limit_val_batches=5,
        num_sanity_val_steps=0,
        max_epochs=4,
        weights_summary=None,
    )
    trainer.fit(model, train_dataloader=train_data, val_dataloaders=val_data)
    trainer.test(model, test_dataloaders=test_data)

if __name__ == '__main__':
    run()

[awaelchli]

Can you describe what the bug is here? And what you expect to happen?
The logging to separate plots was a design choice. Logging a metric in val_step is not meaningful because there is no ordering.

[ethanwharris]

@awaelchli There are a few issues with the current behaviour:

• If you are training for a long time (100 - 200 epochs would be typical), then this creates an unusable number of graphs.
• The original intention (per a comment in the code) was that each epoch would overlay on a single graph, but that is not the case and not easy to do.
• Also, how this interacts with other loggers is unknown, so it's odd that appending /epoch_* is happening always when it targets a specific feature of tensorboard.

The best solution is that people shouldn't really be logging per val step, but when they do the behaviour should at least be consistent with training.