Skip to content

Classification trainer: restore optimizer/scheduler state on resume - #65

Open
jeetv wants to merge 1 commit into
OpenSportsLab:devfrom
mohsayed27:fix/classification-resume-optimizer-state
Open

Classification trainer: restore optimizer/scheduler state on resume#65
jeetv wants to merge 1 commit into
OpenSportsLab:devfrom
mohsayed27:fix/classification-resume-optimizer-state

Conversation

@jeetv

@jeetv jeetv commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator

Trainer_Classification.train() always called build_optimizer()/build_scheduler() fresh, so Adam's momentum and the scheduler's LR-schedule state silently reset on every resume even though _save_checkpoint() had been writing them into every checkpoint all along -- only the epoch counter was ever restored.

Port localization_trainer.py's build_trainer(..., resume_from=...) pattern: Trainer_Classification.load() gains resume_training=True, which builds a fresh optimizer/scheduler from config and restores their state via the existing load_checkpoint() plumbing, stashing the result in self._resume_state. Trainer_Classification.train() gains a resume_from parameter that reuses that state instead of rebuilding it, seeds best-metric tracking, continues from the saved epoch, and guards against resuming a checkpoint that already reached TRAIN.epochs. apis/classification.py wires weights= on ClassificationModel.train() to request resume_training the same way LocalizationModel.train() already does for weights= on its own train().

No behavior change for fresh training (resume_from=None keeps the original build_optimizer()/build_scheduler() path byte-for-byte) or for load_weights() /infer() (resume_training defaults to False there too).

Adds tests/test_classification_trainer_resume.py, which trains a real epoch, saves, and asserts the resumed Adam/StepLR state_dict matches the checkpoint rather than a freshly-initialized optimizer.

Trainer_Classification.train() always called build_optimizer()/build_scheduler()
fresh, so Adam's momentum and the scheduler's LR-schedule state silently reset
on every resume even though _save_checkpoint() had been writing them into
every checkpoint all along -- only the epoch counter was ever restored.

Port localization_trainer.py's build_trainer(..., resume_from=...) pattern:
Trainer_Classification.load() gains resume_training=True, which builds a
fresh optimizer/scheduler from config and restores their state via the
existing load_checkpoint() plumbing, stashing the result in
self._resume_state. Trainer_Classification.train() gains a resume_from
parameter that reuses that state instead of rebuilding it, seeds best-metric
tracking, continues from the saved epoch, and guards against resuming a
checkpoint that already reached TRAIN.epochs. apis/classification.py wires
weights= on ClassificationModel.train() to request resume_training the same
way LocalizationModel.train() already does for weights= on its own train().

No behavior change for fresh training (resume_from=None keeps the original
build_optimizer()/build_scheduler() path byte-for-byte) or for load_weights()
/infer() (resume_training defaults to False there too).

Adds tests/test_classification_trainer_resume.py, which trains a real epoch,
saves, and asserts the resumed Adam/StepLR state_dict matches the checkpoint
rather than a freshly-initialized optimizer.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants