Lightning-AI / Lightning-AI/pytorch-lightning
`ckpt_path` in `Trainer` accepts URIs to automatically load checkpoints from remote paths
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
If I set up a Trainer with a `WandbLogger` and set `log_model=True`, I get that my model is saved locally and in a W&B server.
If I want to retrieve the model from the server I have to use the W&B `use_artifact` methods and the download methods to first retrieve the model, as described [here](https://lightning.ai/docs/pytorch/stable/extensions/generated/lightning.pytorch.loggers.WandbLogger.html)
What I would like to have instead is to not specify the W&B logic explicitly and download the model by just passing a URI with the remote string, (e.g. `wandb://user/project/model-run_id:version`) as the `ckpt_path` in the `Trainer`, along with the `WandbLogger`.
In short, I would like to do something like that:
```
trainer = Trainer(logger=my_wandb_logger, "wanb://user/project/model-run_id:version")
```
### Pitch
Specify a URI to a remote resource, automatically download it, and load it in the `Trainer`.
### Alternatives
Write explicitly the W&B logic to download the model and then pass the checkpoint path to the `Trainer`.
### Additional context
_No response_
cc @lantiga @borda @morganmcg1 @borisdayma @scottire @parambharat
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
The issue names Trainer and WandbLogger as the entry points; start by tracing how Trainer currently handles ckpt_path and how WandbLogger represents logged checkpoints. Define the supported remote URI behavior, then verify that a checkpoint can be downloaded and loaded without caller-managed W&B artifact calls.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 32/100