Lightning-AI/pytorch-lightning

Load callback states while testing.

Open

#5,542 opened on Jan 16, 2021

View on GitHub
 (21 comments) (0 reactions) (0 assignees)Python (3,233 forks)batch import
checkpointingfeaturehelp wantedpriority: 1trainer: testtrainer: validate

Repository metrics

Stars
 (26,687 stars)
PR merge metrics
 (Avg merge 9d 15h) (3 merged PRs in 30d)

Description

🚀 Feature

Load callback states while testing.

Motivation

https://github.com/PyTorchLightning/pytorch-lightning/pull/5161#issuecomment-759083492

Pitch

Two possible API changes:

with an additional argument restore_states:

test(ckpt_path, restore_states=True/False)  # give an option whether to load states or not
test(model, ckpt_path, restore_states=True/False)  # same as above but will just load checkpoint states and not the model

# raise an error
test(ckpt_path=None, restore_states=True)

or without any additional argument:

test(ckpt_path)  # always load states
test(ckpt_path=None)  # don't load any states.
test(model, ckpt_path)  # reload checkpoint states only from ckpt_path

Alternatives

Alternatively, one can just reload checkpoints manually, call on_load_checkpoint for all the callbacks manually, and test.

PS: There may be a better solution. Open to suggestions :) cc: @ananthsub

cc @borda @awaelchli @ananthsub @ninginthecloud @rohitgr7 @tchaton @akihironitta

Contributor guide