Repository navigation
Load the SAM 3.1 multiplex checkpoint once, without spurious key warnings (fixes #506) - #643
Open
YoussefMadkour wants to merge 1 commit into
Open
YoussefMadkour wants to merge 1 commit into
YoussefMadkour wants to merge 1 commit into
Conversation
…ings (facebookresearch#506) build_sam3_multiplex_video_predictor passed the merged checkpoint to build_sam3_multiplex_video_model, which loaded it into the bare tracker. None of the merged keys (tracker.model.*, detector.*) match the tracker, so nothing was loaded and hundreds of missing/unexpected keys were printed; the checkpoint was then read again and loaded correctly into the full demo model. Skip the first load. Also register the real-valued RoPE buffers (freqs_cis_real/imag) as non-persistent: they are derived from freqs_cis at init and are not in the checkpoint, which caused the remaining 64 missing-key warnings.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #506
Problem
Building the SAM 3.1 predictor with
build_sam3_multiplex_video_predictorprints hundreds ofMissing keys/Unexpected keyswarnings, which makes it look like the checkpoint failed to load. There are two causes:checkpoint_pathtobuild_sam3_multiplex_video_model, which loads the merged checkpoint into the bare tracker. The merged checkpoint's keys are prefixedtracker.model.*/detector.*, while the tracker's own keys are not, so this first load matches nothing. It only prints warnings and reads the multi-GB file an extra time. The same file is then loaded again, correctly, into the full demo model a few lines later.use_rope_real=True(the predictor's default), each ViT attention block registersfreqs_cis_real/freqs_cis_imag. These are derived fromfreqs_cisat init and aren't in the checkpoint, which accounts for theMissing keys (64)in the issue.Fix
checkpoint_path=Noneinsidebuild_sam3_multiplex_video_predictor. Its weights still come from the single load into the demo model.freqs_cis_real/freqs_cis_imagwithpersistent=False. They're still created and moved with the module as before, but they're no longer expected in, or saved to, state dicts.freqs_cisitself is unchanged.The weights the model ends up with don't change: I built the tracker on its own and confirmed none of its 931 state-dict keys start with
tracker.ordetector., so the first load never loaded anything.Tests
Added
test/test_multiplex_checkpoint_loading.py(no checkpoint or GPU needed):build_sam3_multiplex_video_predictorno longer passes the checkpoint to the tracker build (it stops the build at that call with a mock).use_rope_real=Truekeeps the real/imag buffers on the module, leaves them out ofstate_dict(), and strictly loads a state dict saved without them.Both tests fail on
mainand pass with this change.