Lightning-AI / Lightning-AI/pytorch-lightning
Allow checkpointing initial model weights
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
Sometimes it can be useful to have access to the randomly initialized weights of a model prior to training. For instance, when evaluating against downstream tasks the untrained model can serve as a random baseline for comparison, especially since many publications are showing that the architectural prior of a randomly initialized model exhibits representational structures that can transfer to downstream applications without any training.
The current `ModelCheckpoint` implementation does not offer an option to allow saving the initial state of a model before any optimization is performed.
### Pitch
The `ModelCheckpoint` callback could be extended by a `save_initial_weights` option that would control whether the initial model is checkpointed prior to training. The flag would default to `False` to prevent breaking behaviour.
### Alternatives
* **Random state can be fixed.** However, this influences other components than the model initialization as well and prevents reproducing a certain model state by simply loading a certain checkpoint.
* **Custom subclass of `ModelCheckpoint`.** Works of course, but would be nice to see this feature as part of PyTorch Lightning.
### Additional context
_No response_
cc @borda @awaelchli @carmocca
Contributor guide
Research direction
Start by locating the ModelCheckpoint callback and tracing when it first saves during training. Add a save_initial_weights option that defaults to False and controls whether the model's pre-training state is checkpointed. Done means the default behavior is unchanged and enabling the option produces an initial checkpoint before optimization.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100