Lightning-AI / Lightning-AI/pytorch-lightning

`ckpt_path` in `Trainer` accepts URIs to automatically load checkpoints from remote paths

Open
#19,849 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature logger: wandb
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.