HybridEmbed for SwinTransformer

[HuKaiLang]

Hi, I'm attempting to use EfficientNet as the backbone network in SwinTransformer, replacing the default PatchEmbed layer.
I want to use EfficientNet's feature extraction capabilities and feed its output into SwinTransformer:

from timm import create_model
from timm.models.vision_transformer_hybrid import HybridEmbed

model = create_model('swin_base_patch4_window7_224', pretrained=True, num_classes=0, in_chans=3)

embedder = create_model('tf_efficientnet_b3_ns', pretrained=True, in_chans=3, features_only=True, out_indices=[2])

hybridembed = HybridEmbed(embedder,img_size=384, 
                                              patch_size=1, 
                                              feature_size=model.patch_embed.grid_size, 
                                              in_chans=3, 
                                              embed_dim=model.embed_dim
                                              )
model.patch_embed = hybridembed

However, SwinTransformer seems not support HybridEmbed, not like convit:

from torchinfo import summary

summary(model, input_size=(8, 3, 224, 224))

File [~/miniconda3/envs/baichuan/lib/python3.10/site-packages/timm/models/swin_transformer.py:322], in SwinTransformerBlock.forward(self, x)
    [321] def forward(self, x):
--> [322]     B, H, W, C = x.shape
    [323]     x = x + self.drop_path1(self._attn(self.norm1(x)))

ValueError: not enough values to unpack (expected 4, got 3)

The code is from here. It seems that SwinTransformer used to support the use of HybridEmbed as patch Embed, but now it doesn't.
Is there another way to implement it?

[rwightman]

@HuKaiLang try adding output_fmt='NHWC' to the HybridEmbed args, believe I changed swin to keep tensors in no-flatten form so they'd be more useful for feature extraction.

[HuKaiLang]

Thank you for your prompt response!

But even though I added output_fmt='NHWC', there was still an additional error:

File [~/miniconda3/envs/baichuan/lib/python3.10/site-packages/timm/models/swin_transformer.py:354], in PatchMerging.forward(self, x)
    [353] B, H, W, C = x.shape
--> [354] _assert(H % 2 == 0, f"x height ({H}) is not even.")
    [355] _assert(W % 2 == 0, f"x width ({W}) is not even.")

File [~/miniconda3/envs/baichuan/lib/python3.10/site-packages/torch/__init__.py:1499], in _assert(condition, message)
   [1498]     return handle_torch_function(_assert, (condition,), condition, message)
-> [1499] assert condition, message

AssertionError: x height (7) is not even.

The above exception was the direct cause of the following exception:

I think that these issues may be related to changes in previous versions of swin_transformer.py.

Since embed_layer appears as a parameter in SwinTransformer's __init__() method, swin_transformer no longer seems to support HybridEmbed.

[rwightman]

@HuKaiLang looks like you just ran out of resolution there, if you have a backbone in the front it will downscale more than the usual swin so you need a higher input res or you need to adjust the transformer block/stride settings as you cannot hit to a non-even input before the last patch merge....

[HuKaiLang]

@rwightman Sorry for the bothering you several times.

Thank you very much for your reply. I tried using a higher resolution input_size, such as 1200 * 1200 and change the img_size in HybridEmbed, but the same code still produces errors:

...

hybridembed = HybridEmbed(embedder,img_size=1200, 
                                              patch_size=1, 
                                              feature_size=model.patch_embed.grid_size, 
                                              in_chans=3, 
                                              embed_dim=model.embed_dim,
                                              output_fmt="NHWC"
                                              )
...

summary(model, input_size=(8, 3, 1200, 1200))

File [~/miniconda3/envs/baichuan/lib/python3.10/site-packages/timm/models/swin_transformer.py:175], in WindowAttention.forward(self, x, mask)
    [174] num_win = mask.shape[0]
--> [175] attn = attn.view(-1, num_win, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
    [176] attn = attn.view(-1, self.num_heads, N, N)

RuntimeError: shape '[-1, 64, 4, 49, 49]' is invalid for input of size 37186688

Although this error occurred in WindowAttention, I think it may be related to the norm1 in the SwinTransformerBlock, and can be traced back to the norm_layer parameter in SwinTransformer. This parameter is also related to embed_layer.

I am not certain whether I can use HybridEmbed on Swin in the same way as ConViT.

[rwightman]

@HuKaiLang you need to spend some time working through the resolutions, window sizing, strides (taking into consideration the additional backbone embed) to get it right, they all need to 'fit' for it to work properly.

[HuKaiLang]

Oh, that's right, it was my mistake, I mistakenly thought there was a problem with the implementation of the Swin.
I checked the code and model parameters, and when I changed the conv_stem.stride to (1,1), the model was able to run.

Thank you very much! I benefited a lot with your help.