Enable to return outputs from {training,validation,test}_step in Manual Optimization 🙏

[jhauret]

Description & Motivation

In previous Lightning versions, it was possible to use multiple optimizers with automatic optimization. Since lightning>=2.0.0 this is no longer possible (i.e. the argument optimizer_idx of the training step disappeared). The problem is not that we have to do the optimization manually, but that we cannot return anything from the {training,validation,test}_step method. Indeed, I have been using the dict return type a lot to put the logging logic inside Lightning hooks like on_train_batch_end which takes the outputs of training_step as input.

Pitch

I would like to be able to return a dict without the key 'loss' from {training,validation,test}_step when using manual optimization.

Alternatives

The only alternative I can find is to move my logging code that is written in on_train_batch_end, on_validation_batch_end, on_test_batch_end of an AbstractLightningModule into the GANLightningModule. The issue is that it generates duplicate code because I have other LightningModules that also inherit from AbstractLightningModule and use the same logging hooks.

Additional context

PS: Is there a typo in this line of the documentation ? Shouldn't it be " Skip to the next batch. This is only supported for automatic manual optimization."

cc @Borda

[awaelchli]

@jhauret Support for returning objects from the training_step has been removed. Previously, the trainer would accumulate these outputs and pass them to the training_epoch_end hooks, but this caused a lot of users shooting themselves in the foot as memory accumulated and led to out-of-memory. It was decided that the user should explicitly collect their objects and handle them so that the implications on memory are clear from the start. You can read the full explanation for this change here in the release notes:

https://github.com/Lightning-AI/pytorch-lightning/releases/tag/2.0.0
Section: The *_epoch_end hooks were removed

There, you also find the pattern how you would do it today in 2.0+. If you just want to log stuff, you can use self.log directly.

The issue is that it generates duplicate code because I have other LightningModules that also inherit from AbstractLightningModule and use the same logging hooks.

You can avoid code duplication through composition or refactoring it into a function for example.

PS: Is there a typo in this line of the documentation ? Shouldn't it be " Skip to the next batch. This is only supported for automatic manual optimization."

No I'm sure this is correct. Skipping the training_step by returning None means no optimization will be done, and the loop continues to the next batch. This makes sense only for automatic optimization.

[jhauret]

Thanks for your reply @awaelchli !

I agree that the accumulation of outputs for the former training_epoch_end(self, outputs) was dangerous, and you made a good point in removing it in 2.0.+. My request was just to be able to use the on_train_batch_end(outputs, batch, batch_idx), which still exists, but in the manual optimization scheme. (not forced to return None in training_step). Still, your explanations + the link you provided really helped me to see what design I should aim for!

And for the last point, since the optimization is done inside the training_step function in the manual optimization. I think returning None is supported for both automatic and manual optimization scenarios. In fact, from my understanding of the current version, it is even mandatory for the manual optimization mode.

[awaelchli]

Hey
Sorry for the misunderstanding. Yes I acknowledge that you asked for on_train_batch_end, apologies!

This is supported, here is an example where I use manual optimization, return a dummy value from training_step, and print it in on_train_batch_end.

import torch
from lightning.pytorch import LightningModule, Trainer
from torch.utils.data import DataLoader, Dataset


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]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)
        self.automatic_optimization = False

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

    def training_step(self, batch, batch_idx):
        optimizer = self.optimizers()
        loss = self(batch).sum()
        self.backward(loss)
        optimizer.step()
        return {"batch_idx": batch_idx * 2}

    def on_train_batch_end(self, outputs, batch, batch_idx):
        print(outputs)

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


train_data = DataLoader(RandomDataset(32, 64), batch_size=2)

model = BoringModel()
trainer = Trainer(max_epochs=1, enable_progress_bar=False)
trainer.fit(model, train_dataloaders=train_data)

Is this what you are looking for?

In fact, from my understanding of the current version, it is even mandatory for the manual optimization mode.

Also here a quick "proof" that returning something in manual optimization is not required:

import torch
from lightning.pytorch import LightningModule, Trainer
from torch.utils.data import DataLoader, Dataset


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]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)
        self.automatic_optimization = False

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

    def training_step(self, batch, batch_idx):
        optimizer = self.optimizers()
        loss = self(batch).sum()
        self.backward(loss)
        optimizer.step()
        return None  # This is allowed, not required!

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


train_data = DataLoader(RandomDataset(32, 64), batch_size=2)

model = BoringModel()
trainer = Trainer(max_epochs=1, enable_progress_bar=True)
trainer.fit(model, train_dataloaders=train_data)

[jhauret]

I am happy to see that what I wanted is already possible then ! 🥳

For the sake of clarity, maybe you could just update the documentation of the training_step Returns. Indeed the training_step of the first BoringModel is not one of the three options proposed in the doc (see below) and the second BoringModel's training_step output is None, so it is also supported for manual optimization.

---

Documentation says:

• Tensor - The loss tensor
• dict - A dictionary. Can include any keys, but must include the key 'loss'.
• None - Skip to the next batch. This is only supported for automatic optimization.

---

Thanks again for your answers and your patience!

[awaelchli]

Yes we can update the text in the docs. Does the way it is described in the linked PR work better now?

[jhauret]

It's perfect like that 👌