Lightning-AI / Lightning-AI/pytorch-lightning
Add `after_instantiate_classes` hook to LightningCLI to save LightningDataModule information in the current log directory
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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