Lightning-AI / Lightning-AI/pytorch-lightning

Add `after_instantiate_classes` hook to LightningCLI to save LightningDataModule information in the current log directory

Open
#18,608 7 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature lightningcli
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Description & Motivation

Hi.

I noticed that `LightningCLI` has a member function called `before_instantiate_classes` called before `instantiate_classes` but no `after_instantiate_classes`. I need to save the data split generated in `LightningDataModule` after it has been instantiated to a file in the current logging directory. And this is only possible after `Trainer.instantiate_classes`, because only then I have access to `trainer.logger.log_dir` and `LightningCLI.datamodule` in order to save files related to the current run in the current logging directory.

I tried to do it with a custom Callback for the Trainer overwriting `on_fit_start()`, however, the trainer only has access to the dataloaders, and not the underlying datamodule so you can't do it that way.

In my specific usecase, I want to save and load the training/validation/test data split to and from a separate file based on the result of `torch.utils.data.random_split` or if the filepath are given, based on the given files. The goal is to train different models (in different LightningCLI instances) on the same data split given the same underlying data, without doing requiring setting any seed, because there can be a case where you only pick some of the data to run validation on, and not all of them.
```python
# In the LightningDataModule
data = #Inherits for example from torch DatasetFolder

if data_split:
train_data, val_data = random_split(data, data_split)
else:
train-data = Subset(data, train_indices_from_file)
val_data-data = Subset(data, val_indices_from_file)
```

During training I could basically save the `state_dict` of the `LightningDataModule` in the current checkpoint, however, this will incur additional complexity in extracting the information about the data splits whenever I want to train another model with LightningCLI on the same splits.

For now, let's imagine the locally stored dataset does not change and I don't have to worry about the stored data.

My solution now is just overriding the `instantiate_classes` member function as follows:
```python
    class MyLightningCLI(LightningCLI):
        def instantiate_classes(self) -> None:
            """Instantiates the classes and sets their attributes."""
            self.config_init = self.parser.instantiate_classes(self.config)
            self.datamodule = self._get(self.config_init, "data")
            self.model = self._get(self.config_init, "model")
            self._add_configure_optimizers_method_to_model(self.subcommand)
            self.trainer = self.instantiate_trainer()

# this is the new change
            self.datamodule.save_data_split_to_file(self.trainer.logger.log_dir)
```
where `datamodule.save_data_split_to_file` is just a member of the LightningDataModule class which implements the storing functionality. I also had to move the creation of the train and val data splits from `setup()` to the constructor of this class because otherwise I wouldn't have the information about the splits generated when calling `MyLightningCLI.instantiate_classes`

If we could add another hook called `after_instantiate_classes` called after `LightningCLI.instantiate_classes` then one wouldn't have to override `instantiate_classes` and the code would look more "lightning-ish".

Thank you

### Pitch

Add `after_instantiate_classes` hook to be called after `instantiate_classes` in LightningCLI.

### Alternatives

_No response_

### Additional context

_No response_

cc @borda @carmocca @mauvilsa

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with LightningCLI.instantiate_classes and the existing before_instantiate_classes hook described in the issue. Trace when datamodule and trainer become available, then add and exercise the post-instantiation hook so subclasses can access the datamodule and trainer logger directory after instantiation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
cli, machine-learning
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.