Skip to content

Fix: Checkpoint saving and loading#370

Merged
kozlov721 merged 30 commits intomainfrom
fix/custom-checkpoint-callback
Apr 15, 2026
Merged

Fix: Checkpoint saving and loading#370
kozlov721 merged 30 commits intomainfrom
fix/custom-checkpoint-callback

Conversation

@kozlov721
Copy link
Copy Markdown
Collaborator

@kozlov721 kozlov721 commented Apr 10, 2026 •
edited
Loading

Purpose

Fixes a few issues with checkpoint saving and loading:

  • The UploadCheckpoint.on_save_checkpoint method is called before the checkpoint is modified by the LuxonisLightningModule.on_save_checkpoint
    • Causes the uploaded checkpoint to lack the extra checkpoint fields
  • Inconsistent loading of dataset metadata from a checkpoint

Specification

  • Fixed adding our extra keys to the checkpoint in UploadCheckpoint
  • Added save_checkpoint and get_checkpoint methods to LuxonisModel for easier testing
  • Moved config restoration from checkpoint from __main__ to LuxonisModel
  • Renamed DebugLoader to DummyLoader and debug_mode in LuxonisModel to allow_empty_dataset
  • Added class_names parameter to DummyLoader
  • Renamed debug_mode to allow_empty_dataset in LuxonisModel

Dependencies & Potential Impact

None / not applicable

Deployment Plan

None / not applicable

Testing & Validation

  • Added new checkpoint tests

@kozlov721 kozlov721 requested a review from a team as a code owner April 10, 2026 23:52
@kozlov721 kozlov721 requested review from conorsim, klemen1999 and tersekmatija and removed request for a team April 10, 2026 23:52
@github-actions github-actions bot added the fix Fixing a bug label Apr 10, 2026
@kozlov721 kozlov721 force-pushed the fix/custom-checkpoint-callback branch from c8a16bf to 3e98228 Compare April 13, 2026 12:20
@github-actions github-actions bot added the CLI Changes affecting the CLI label Apr 13, 2026
@kozlov721 kozlov721 force-pushed the fix/custom-checkpoint-callback branch from 9b1d05d to f3c7c92 Compare April 13, 2026 14:30
@kozlov721 kozlov721 changed the title Fix: Save custom checkpoint info in UploadCheckpoint Fix: Checkpoint saving and loading Apr 13, 2026
klemen1999
klemen1999 reviewed Apr 15, 2026
View reviewed changes
Comment thread luxonis_train/__main__.py
@github-actions github-actions bot added the documentation Improvements or additions to documentation label Apr 15, 2026
@kozlov721 kozlov721 merged commit 4c125c1 into main Apr 15, 2026
11 checks passed
@kozlov721 kozlov721 deleted the fix/custom-checkpoint-callback branch April 15, 2026 12:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@klemen1999 klemen1999 klemen1999 approved these changes

@tersekmatija tersekmatija Awaiting requested review from tersekmatija tersekmatija is a code owner automatically assigned from luxonis/ml-reviewers

@conorsim conorsim Awaiting requested review from conorsim conorsim is a code owner automatically assigned from luxonis/ml-reviewers

Assignees

@kozlov721 kozlov721

Labels

CLI Changes affecting the CLI documentation Improvements or additions to documentation fix Fixing a bug

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@kozlov721 @klemen1999