Lightning-AI / Lightning-AI/pytorch-lightning

Allow checkpointing initial model weights

Open
#17,469 3 comments 0 reactions 0 assignees View on GitHub
callback: model checkpoint checkpointing feature
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.