Forward function. Args: x: Input feature, tensor size (B, H*W, C). H, W: Spatial resolution of the input feature.
(self, x, H, W)
| 259 | self.norm = norm_layer(4 * dim) |
| 260 | |
| 261 | def forward(self, x, H, W): |
| 262 | """ Forward function. |
| 263 | |
| 264 | Args: |
| 265 | x: Input feature, tensor size (B, H*W, C). |
| 266 | H, W: Spatial resolution of the input feature. |
| 267 | """ |
| 268 | B, L, C = x.shape |
| 269 | assert L == H * W, "input feature has wrong size" |
| 270 | |
| 271 | x = x.view(B, H, W, C) |
| 272 | |
| 273 | # padding |
| 274 | pad_input = (H % 2 == 1) or (W % 2 == 1) |
| 275 | if pad_input: |
| 276 | x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) |
| 277 | |
| 278 | x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C |
| 279 | x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C |
| 280 | x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C |
| 281 | x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C |
| 282 | x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C |
| 283 | x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C |
| 284 | |
| 285 | x = self.norm(x) |
| 286 | x = self.reduction(x) |
| 287 | |
| 288 | return x |
| 289 | |
| 290 | |
| 291 | class PatchEmbed(nn.Module): |