Lightning-AI / Lightning-AI/pytorch-lightning
LightningCLI feedback: improvement suggestions
@edenlightning is already working on this.
Since Apr 17, 2023.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
# 🚀 Feature
## Motivation
I've started using (i.e. subclassing) the `LightningCLI` recently and, given that it is still in beta, I decided to share with you some patterns that I end up using in every one of my projects. If other people also rely on these patterns often, it might be valuable to include them in the base class, smoothing out its usage, ensuring that it "just works" out of the box.
Let me know what you think!
## Disclaimer
I do not use the CLI for model tuning, so my feedback is solely focused on the other commands: `fit`, `validate`, `test` and `predict`.
## Pitch
For ease of exposition and discussion, I will show each "pattern" as a separate entity, although each of them ends up modifying the same few `LightningCLI` methods and thus my overall code looks slightly different than the simple concatenation of these examples.
### 1) Experiment folder
I typically want each run to have its own dedicated folder for checkpoints, logs and other artifacts. If I'm running fit, chances are I want this directory to be created for me. If I'm doing anything else, instead, I'm happy to get an error if it does not exist.
```python
def add_arguments_to_parser(self, parser):
parser.add_argument('--directory', type = str)
# other things here
def before_instantiate_classes(self):
# other things here
if subcommand == 'fit':
os.makedirs(self.config['fit.directory'], exist_ok = True)
# other things here
```
Right now, I explicitly use this folder in my yaml files to setup the trainer and each individual logger, but that's something that I should probably automate in the Python code as well, to avoid mistakes and repetitions in the config file:
```yaml
name:
directory: runs/${name}
trainer:
default_root_dir: ${directory}
logger:
- class_path: pytorch_lightning.loggers.WandbLogger
init_args:
save_dir: ${directory}
- class_path: pytorch_lightning.loggers.CSVLogger
init_args:
save_dir: ${directory}
```
So I should probably do it as follows:
```python
# WARNING: this particular snippet was not tested
def before_instantiate_classes(self):
# other things here
directory = self.config[f'{subcommand}.directory']
self.config[f'{subcommand}.trainer.default_root_dir'] = directory
for logger_conf in self.config[f'{subcommand}.trainer.logger']:
logger_conf['init_args']['save_dir'] = directory
# other things here
```
### 2) Fit-only behaviour
Creating the experiment folder is not the only thing that I only want to do during training. `validate`, `test` and `predict` are typically one-off things that I do once I'm happy with the fit behaviour, which I check from the logs.
As such, when I run one of these other three commands on an existing experiment, as follows:
```bash
python -m scripts.cli --config "runs//config.yaml"
```
I typically have the following requirements:
1) I don't want to overwrite that configuration file if I pass any extra args, such as `--trainer.gpus`, because I want that configuration file to reflect the way that the model was trained during fit.
2) I don't want to overwrite the local log files or upload a new run to online logging platforms, as I'm not actually doing a new fit.
To achieve these, I end up doing the following:
```python
def before_instantiate_classes(self):
# other things here
if subcommand != 'fit':
self.save_config_callback = None
self.config[f'{subcommand}.trainer.logger'] = False
# other things here
```
### 3) Automated retrieval of best checkpoint
Again, when I'm using validate/test/predict, unless explicitly specified, I typically want to perform that command on the best checkpoint from the fit phase.
The `Trainer` class has a `ckpt_path='best'` option to achieve this, but it only works if that trainer is the same that was used during fit. Unfortunately, this is not the case when calling validate/test/predict from the command line.
So I end up also adding the following code:
```python
def before_instantiate_classes(self):
# other things here
if subcommand != 'fit':
self.config[f'{subcommand}.ckpt_path'] = os.path.join(self.config[f'{subcommand}.directory'], 'best.ckpt')
# other things here
def fit(self, **trainer_kwargs):
self.trainer.fit(**trainer_kwargs)
os.symlink(self.trainer.checkpoint_callback.best_model_path, os.path.join(self.config['fit.directory'], 'best.ckpt'))
```
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.
Assessment
This issue has not been assessed yet.