x: B, H*W, C
(self, x, size)
| 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 | |
| 719 | def img2windows(img, H_sp, W_sp): |
| 720 | """ |
nothing calls this directly
no outgoing calls
no test coverage detected