| 332 | |
| 333 | |
| 334 | class MultiModalSwinTransformer(nn.Module): |
| 335 | def __init__(self, |
| 336 | pretrain_img_size=224, |
| 337 | patch_size=4, |
| 338 | in_chans=3, |
| 339 | embed_dim=96, |
| 340 | depths=[2, 2, 6, 2], |
| 341 | num_heads=[3, 6, 12, 24], |
| 342 | window_size=7, |
| 343 | mlp_ratio=4., |
| 344 | qkv_bias=True, |
| 345 | qk_scale=None, |
| 346 | drop_rate=0., |
| 347 | attn_drop_rate=0., |
| 348 | drop_path_rate=0.2, |
| 349 | norm_layer=nn.LayerNorm, |
| 350 | ape=False, |
| 351 | patch_norm=True, |
| 352 | out_indices=(0, 1, 2, 3), |
| 353 | frozen_stages=-1, |
| 354 | use_checkpoint=False, |
| 355 | num_heads_fusion=[1, 1, 1, 1], |
| 356 | fusion_drop=0.0 |
| 357 | ): |
| 358 | super().__init__() |
| 359 | |
| 360 | self.pretrain_img_size = pretrain_img_size |
| 361 | self.num_layers = len(depths) |
| 362 | self.embed_dim = embed_dim |
| 363 | self.ape = ape |
| 364 | self.patch_norm = patch_norm |
| 365 | self.out_indices = out_indices |
| 366 | self.frozen_stages = frozen_stages |
| 367 | |
| 368 | # split image into non-overlapping patches |
| 369 | self.patch_embed = PatchEmbed( |
| 370 | patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, |
| 371 | norm_layer=norm_layer if self.patch_norm else None) |
| 372 | |
| 373 | # absolute position embedding |
| 374 | if self.ape: |
| 375 | pretrain_img_size = to_2tuple(pretrain_img_size) |
| 376 | patch_size = to_2tuple(patch_size) |
| 377 | patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]] |
| 378 | |
| 379 | self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])) |
| 380 | trunc_normal_(self.absolute_pos_embed, std=.02) |
| 381 | |
| 382 | self.pos_drop = nn.Dropout(p=drop_rate) |
| 383 | |
| 384 | # stochastic depth |
| 385 | dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule |
| 386 | |
| 387 | # build layers |
| 388 | self.layers = nn.ModuleList() |
| 389 | for i_layer in range(self.num_layers): |
| 390 | layer = MMBasicLayer( |
| 391 | dim=int(embed_dim * 2 ** i_layer), |