[FEATURE] Segment Anything backbones

[SimJeg]

Hello,

Maybe it's out of the scope of timm so far, but it could be interestint to add support for the 3 models released by Meta (ViT-B/L/H) in their Segment Anything project. These models are pretrained using MAE and then finetuned on 11M images for segmentation. An interesting feature is that these models have been trained with an image size of 1024x1024. The ViT are implemented in pure pytorch here without using scaled_dot_product_attention. Curious to see how these models transfer to classification tasks :)

Thanks,
Simon

[leng-yue]

It seems like we need to add the window attention block.

[seefun]

It seems like we need to add the window attention block.

It seems that ViT from SAM has the same arch with ViTDet; There are many differences from the Vanilla ViT, exclude window attention and different pos encoding method; I'll give it a try.

[seefun]

Due to the non standard ViT structure used by Segment-Anything Model(SAM)：

• window attention just like swin-transformer (no-overlap)
• relative position encoding just like vision_transformer_relpos
• some other different details like global attention and windows attention choose

So I created a new class for SAM's ViT, and load the pretrained to check whether the pre training weight loaded correctly.

Some changes may be needed, such as supporting any shape input and change the code style. I'm not sure if my code is possible to provide a reference for you.

#1812

[rwightman]

@SimJeg segment anything backbone is in thanks to @seefun, I did a last min tweak to make the resolution dynamic (as long as it's modulus patch_size), might make that change for other vits if it works w/o issues. Currently w/o F.sdpa as it needs more work due to the way the rel pos is applied

[tommiekerssies]

I consistently find worse fine-tuning performance from timm's SAM implementation versus the original (with all else the same). Any clues on why this may be?

[seefun]

I consistently find worse fine-tuning performance from timm's SAM implementation versus the original (with all else the same). Any clues on why this may be?

For Deep Learning Network, In theory, as long as the forward result is consistent, the gradient of image is also consistent. But the gradient of each parameter in the network may not be consistent. So you can try to observe the gradient of each parameter and check if it is consistent

Are there some experiment results and more details?

Some parameters should be set to no weight decay when finetuning: 'pos_embed', 'dist_token';
In timm's impl, there are something like droppath, default set to '0.0', It may also be the reason for the differences

[rwightman]

@tommiekerssies @seefun I found a bug that I introduced in the q scaling wrt to rel pos when I reorganized things towards a more typical form I use in timm (for future F.sdpa support). I've now finished adding the F.sdpa support and fixed that issue so hopefully fine-tune is better #1920

[tommiekerssies]

@rwightman great work! They now have similar performance (the original vs timm implementation). One note here is that if I do not want to use the neck, I have to manually replace it by a LayerNorm. If I set out_chans=None, nothing happens. If I set out_chans=0, I will get an error when loading the weights.