Allow epochs with different size

[lucian-student]

Description & Motivation

Currently there is no support for epochs of different lengths.
The reason is for example what if I want use augmentation that makes dataset shorter.

For example my usacase is to offset timeseries data in each epoch differently.

Pitch

For timeseries data it would be very useful if logging didnt break, when there are epochs of different lengths.

So far logging completely breaks when u dynamically change epoch length.

Alternatives

Sure i could make it work by merging everything into one epoch. But that would disallow me to use on-epoch hooks, which I want to utilize.

So far Iam going to inspect logging logic, because it seems there is an issue. Or hook trigger logic, since so far iam not sure if hooks even trigger for epoch end.

Additional context

No response

cc @lantiga

[deependujha]

Thanks for the report @lucian-student!

This is actually working as designed. By default, Lightning only creates the training dataloader once (at the start of training) and reuses it for all epochs.

If your dataset changes size between epochs, you need to tell the Trainer to reload the dataloader:

pytorch-lightning/src/lightning/pytorch/trainer/trainer.py

Lines 300 to 301 in fe6b1cc

	reload_dataloaders_every_n_epochs: Set to a positive integer to reload dataloaders every n epochs.
	Default: ``0``.

trainer = pl.Trainer(reload_dataloaders_every_n_epochs=1)

This ensures that train_dataloader() is called before each epoch, so any changes to dataset size, augmentation offsets, etc., are correctly reflected, including progress bar totals, validation scheduling, and metric logging.

Here's a minimal working example:

import torch
from torch.utils.data import Dataset, DataLoader
import lightning.pytorch as pl


class VariableSizeDataset(Dataset):
    def __init__(self, size):
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        x = torch.randn(4)
        y = torch.tensor(idx % 2, dtype=torch.float32)
        return x, y


class TimeSeriesDataModule(pl.LightningDataModule):
    EPOCH_SIZES = [80, 30, 120]

    def __init__(self, batch_size=10):
        super().__init__()
        self.batch_size = batch_size
        self._call_count = 0

    def train_dataloader(self):
        size = self.EPOCH_SIZES[self._call_count % len(self.EPOCH_SIZES)]
        self._call_count += 1
        return DataLoader(
            VariableSizeDataset(size), batch_size=self.batch_size
        )


class SimpleModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(4, 1)

    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = torch.nn.functional.mse_loss(self.layer(x).squeeze(), y)
        self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True)
        return loss

    def on_train_epoch_end(self):
        print(
            f"Epoch {self.current_epoch}: "
            f"num_training_batches={self.trainer.num_training_batches}"
        )

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


trainer = pl.Trainer(
    max_epochs=3,
    reload_dataloaders_every_n_epochs=1,  # <-- the key setting
    accelerator="cpu",
)
trainer.fit(SimpleModel(), datamodule=TimeSeriesDataModule())

Expected:
    Epoch 0: num_training_batches=8
    Epoch 1: num_training_batches=3
    Epoch 2: num_training_batches=12

Could you verify whether reload_dataloaders_every_n_epochs=1 resolves your issue? If logging still breaks with that setting enabled, please provide a minimal reproducer and we'll investigate further.

[lucian-student]

@deependujha

For me reload_dataloaders_every_n_epochs didn't work, so idk. I might be holding it wrong.

[deependujha]

@lucian-student, can you please share minimal reproducible.

[lucian-student]

@deependujha sec once iam back from work I will circle back to you.