(self, dim, reso, num_heads,
split_size=7, mlp_ratio=4., qkv_bias=False, qk_scale=None,
drop=0., attn_drop=0., drop_path=0.,
act_layer=nn.GELU, norm_layer=nn.LayerNorm,
last_stage=False)
| 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 | """ |
nothing calls this directly
no test coverage detected