Embedding dimension mismatch (expected 128 but got 256) when using UltraSAM.pth with custom pipeline

[priteeharmale13-hash]

I’m encountering an embedding dimension mismatch when using UltraSAM.pth in my custom pipeline, even though the official example runs successfully.

Error: AssertionError: was expecting embedding dimension of 128, but got 256.

I verified the checkpoint embedding size using:
python -c "import torch; print([v.shape for k,v in torch.load('UltraSAM.pth', map_location='cpu')['state_dict'].items() if 'prompt_encoder' in k and len(v.shape)==2][0])"
output:
torch.Size([2, 128])

What is the correct way to configure the model backbone and decoder to use UltraSAM.pth in a custom pipeline (for segmentation)?

Is there any specific modification required in the custom pipeline / data preprocessor for UltraSAM to generate segmentation masks?

[Chop1]

Hi, are you running our code?

The SAM use a modified attention, which is different from the one in pytorch, that reduce the size from 256 to 128.

In our code, we have a hook,

UltraSam/endosam/hooks/MonkeyPatchHook.py

Line 6 in 458ec79

@HOOKS.register_module()

Which is call before use, see example in our config,

UltraSam/configs/_base_/models/sam.py

Line 127 in 458ec79

custom_hooks = [dict(type="MonkeyPatchHook")]

this is to be in adequation with the weight from the official SAM from meta.

I believe this is what is triggering this issue. Make sure to run the hook, or monkeypatch as we did.

[priteeharmale13-hash]

Thanks for the clarification regarding the MonkeyPatchHook. I added the hook as shown in your config (custom_hooks = [dict(type="MonkeyPatchHook")]), and the model now initializes correctly.
However, I’m now encountering a new runtime error during prediction:
DEBUG → GT BBOX shape: (1, 4)
DEBUG → input scale_factor shape: (8,)
Traceback (most recent call last):
File "test_ultrasam_box_segmentation.py", line 97, in
pred = model(**processed, mode="predict")[0]
...
File "mmdet/structures/bbox/transforms.py", line 414, in scale_boxes
return boxes * scale_factor
RuntimeError: The size of tensor a (4) must match the size of tensor b (16) at non-singleton dimension 1
It seems the scale_factor in b.metainfo has a shape mismatch — (8,) instead of (4,).

This happens at this line inside endosam/models/detectors/SAM.py:
pred_box_prompts = scale_boxes(r['bboxes'], b.metainfo['scale_factor']).view(...)
Could you please confirm:

Should scale_factor contain 4 elements (for x, y scaling) or is it expected to have 8 values in the modified pipeline?

[Chop1]

I mean, i don't know what is your modified pipeline and code.
scale_factor (from top of my mind) should be a tuple of 2 values.

I invite you to run our example, and inspect the variables shape and code.