Set the optimization output result class as a class attribute

[carmocca]

What does this PR do?

Allow configuring the output processing class without subclassing

Partly addresses #9968

With this, the following would be possible

import os
from dataclasses import dataclass

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

from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.loops.optimization.optimizer_loop import ClosureResult


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)

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

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return loss

    def training_epoch_end(self, outputs):
        print(outputs)

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


@dataclass
class MyClosureResult(ClosureResult):
    was_dict: bool = False

    @classmethod
    def from_training_step_output(cls, training_step_output, *args, **kwargs):
        super_obj = ClosureResult.from_training_step_output(training_step_output, *args, **kwargs)
        obj = cls(super_obj.closure_loss, super_obj.extra, was_dict=isinstance(training_step_output, dict))
        return obj

    def asdict(self):
        if not self.was_dict:
            return self.loss
        return super().asdict()


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        enable_model_summary=False,
        enable_progress_bar=False,
    )

    trainer.fit_loop.epoch_loop.batch_loop.optimizer_loop.output_result_cls = MyClosureResult

    trainer.fit(model, train_dataloaders=train_data)


if __name__ == "__main__":
    run()

Does your PR introduce any breaking changes? If yes, please list them.

None

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• [n/a] Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

[codecov]

Codecov Report

Merging #9977 (d175a9d) into master (788f686) will increase coverage by 0%.
The diff coverage is 100%.

@@          Coverage Diff           @@
##           master   #9977   +/-   ##
======================================
  Coverage      93%     93%           
======================================
  Files         179     179           
  Lines       15833   15835    +2     
======================================
+ Hits        14673   14675    +2     
  Misses       1160    1160