Skip to content

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

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

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

Conversation

@mohsayed27

Copy link
Copy Markdown

Summary

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 both into every checkpoint all along. Only the epoch counter was ever actually restored.

localization_trainer.py's build_trainer(..., resume_from=...) already solves this correctly for the localization task. This PR ports that pattern to the classification trainer:

  • Trainer_Classification.load() gains resume_training=False; when True, it rebuilds the optimizer/scheduler from config and restores their state via the existing load_checkpoint() plumbing (which already supported this — it just had no caller wiring it up for classification), stashing the result in self._resume_state.
  • Trainer_Classification.train() gains resume_from=None; when given, it reuses that optimizer/scheduler instead of rebuilding them, continues from the saved epoch, seeds best-metric tracking, and raises a clear error if the checkpoint already reached TRAIN.epochs (mirroring localization's guard).
  • apis/classification.py wires weights= on ClassificationModel.train() to request resume_training, the same way LocalizationModel.train() already treats weights= as a resume trigger — no new public API surface.

Fresh training (resume_from=None) is unchanged — verified byte-for-byte identical code path — and so is load_weights()/inference (resume_training defaults to False there too).

Test plan

  • Added tests/test_classification_trainer_resume.py: trains one real epoch with a real Adam optimizer, saves via _save_checkpoint, resumes, and asserts the resumed optimizer's state_dict()["state"] matches the checkpoint's saved momentum buffers field-by-field (while a freshly-built optimizer over the same model has empty state) — proves actual restoration, not just "doesn't crash."
  • Verified the test is meaningful: reverting the source changes makes 3 of the 4 new tests fail with TypeError on the new keyword arguments, confirming they're genuinely exercising the new code.
  • All 4 new tests pass; ran on the current main tip after rebasing.

🤖 Generated with Claude Code

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.
@jeetv jeetv added the invalid This doesn't seem right label Sep 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

invalid This doesn't seem right

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants