Configure model checkpointing - Githubissues

Clay-foundation / model

The Clay Foundation Model (in development)

https://clay-foundation.github.io/model/

Apache License 2.0

362 stars 46 forks source link

Configure model checkpointing #55

Closed weiji14 closed 11 months ago

weiji14 commented 11 months ago

What I am changing

Configuring how model weights are saved to checkpoint (*.ckpt) files during training

How I did it

Using Lightning's ModelCheckpoint callback with the following parameters:
- auto_insert_metric_name=False
- filename="vit_epoch-{epoch:02d}_train_loss-{train/loss:.2f}"
- monitor="train/loss"
- mode="min"
- save_last=True
- save_top_k=3
- save_weights_only=True
Added the AsyncCheckpointIO plugin to reduce the gap in time between epochs when a checkpoint is saved

How you can test it

Run python trainer.py fit --trainer.max_epochs=20 locally
Extra configuration options can be found using python trainer.py fit --trainer.callbacks.help=ModelCheckpoint

Related Issues

References:

https://github.com/Lightning-AI/lightning/issues/11561

weiji14 commented 11 months ago

Gonna merge this in directly since it's a relatively small change, but useful in terms of showing how callbacks and plugins can be added to the LightningCLI in trainer.py.