| 495 | |
| 496 | |
| 497 | class MMBasicLayer(nn.Module): |
| 498 | def __init__(self, |
| 499 | dim, |
| 500 | depth, |
| 501 | num_heads, |
| 502 | window_size=7, |
| 503 | mlp_ratio=4., |
| 504 | qkv_bias=True, |
| 505 | qk_scale=None, |
| 506 | drop=0., |
| 507 | attn_drop=0., |
| 508 | drop_path=0., |
| 509 | norm_layer=nn.LayerNorm, |
| 510 | downsample=None, |
| 511 | use_checkpoint=False, |
| 512 | num_heads_fusion=1, |
| 513 | fusion_drop=0.0 |
| 514 | ): |
| 515 | super().__init__() |
| 516 | self.window_size = window_size |
| 517 | self.shift_size = window_size // 2 |
| 518 | self.depth = depth |
| 519 | self.use_checkpoint = use_checkpoint |
| 520 | self.dim = dim |
| 521 | |
| 522 | # build blocks |
| 523 | self.blocks = nn.ModuleList([ |
| 524 | SwinTransformerBlock( |
| 525 | dim=dim, |
| 526 | num_heads=num_heads, |
| 527 | window_size=window_size, |
| 528 | shift_size=0 if (i % 2 == 0) else window_size // 2, |
| 529 | mlp_ratio=mlp_ratio, |
| 530 | qkv_bias=qkv_bias, |
| 531 | qk_scale=qk_scale, |
| 532 | drop=drop, |
| 533 | attn_drop=attn_drop, |
| 534 | drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, |
| 535 | norm_layer=norm_layer) |
| 536 | for i in range(depth)]) |
| 537 | |
| 538 | # fuse before downsampling |
| 539 | self.pwam_fusion = PWAM(dim, # both the visual input and for combining, num of channels |
| 540 | dim, # v_in |
| 541 | 768, # l_in |
| 542 | dim, # key |
| 543 | dim, # value |
| 544 | num_heads=num_heads_fusion, |
| 545 | dropout=fusion_drop) |
| 546 | |
| 547 | #self.pwam_lattn = SelfAttn(qkv_dim=768, hidden_dim=768, out_dim=768, bias=False, attn_drop=.1, proj_drop=.1) |
| 548 | |
| 549 | self.pwam_gate = nn.Sequential( |
| 550 | nn.Linear(dim, dim, bias=False), |
| 551 | nn.ReLU(), |
| 552 | nn.Linear(dim, dim, bias=False), |
| 553 | nn.Tanh() |
| 554 | ) |