| 659 | |
| 660 | |
| 661 | class CSWinBlock(nn.Module): |
| 662 | def __init__(self, dim, reso, num_heads, |
| 663 | split_size=7, mlp_ratio=4., qkv_bias=False, qk_scale=None, |
| 664 | drop=0., attn_drop=0., drop_path=0., |
| 665 | act_layer=nn.GELU, norm_layer=nn.LayerNorm, |
| 666 | last_stage=False): |
| 667 | super().__init__() |
| 668 | self.dim = dim |
| 669 | self.num_heads = num_heads |
| 670 | self.patches_resolution = reso |
| 671 | self.split_size = split_size |
| 672 | self.mlp_ratio = mlp_ratio |
| 673 | self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) |
| 674 | self.norm1 = norm_layer(dim) |
| 675 | |
| 676 | last_stage = False |
| 677 | if last_stage: |
| 678 | self.branch_num = 1 |
| 679 | else: |
| 680 | self.branch_num = 2 |
| 681 | self.proj = nn.Linear(dim, dim) |
| 682 | self.proj_drop = nn.Dropout(drop) |
| 683 | |
| 684 | self.attns = nn.ModuleList([ |
| 685 | LePEAttention( |
| 686 | dim // 2, resolution=self.patches_resolution, idx=i, |
| 687 | split_size=split_size, num_heads=num_heads // 2, dim_out=dim // 2, |
| 688 | qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) |
| 689 | for i in range(self.branch_num)]) |
| 690 | |
| 691 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 692 | |
| 693 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 694 | self.mlp = FeedForward(in_dim=dim, hidden_dim=mlp_hidden_dim, act_layer=act_layer, dropout=drop) |
| 695 | self.norm2 = norm_layer(dim) |
| 696 | |
| 697 | def forward(self, x, size): |
| 698 | """ |
| 699 | x: B, H*W, C |
| 700 | """ |
| 701 | H, W = size |
| 702 | B, L, C = x.shape |
| 703 | assert L == H * W, "flatten img_tokens has wrong size" |
| 704 | img = self.norm1(x) |
| 705 | qkv = self.qkv(img).reshape(B, -1, 3, C).permute(2, 0, 1, 3) |
| 706 | |
| 707 | if self.branch_num == 2: |
| 708 | x1 = self.attns[0](qkv[:, :, :, :C // 2], size) |
| 709 | x2 = self.attns[1](qkv[:, :, :, C // 2:], size) |
| 710 | attened_x = torch.cat([x1, x2], dim=2) |
| 711 | else: |
| 712 | attened_x = self.attns[0](qkv, size) |
| 713 | attened_x = self.proj(attened_x) |
| 714 | x = x + self.drop_path(attened_x) |
| 715 | x = x + self.drop_path(self.mlp(self.norm2(x))) |
| 716 | |
| 717 | return x |
| 718 | |