Forward function. Args: x: Input feature, tensor size (B, H*W, C). H, W: Spatial resolution of the input feature.
(self, x, H, W)
| 320 | self.norm = norm_layer(4 * dim) |
| 321 | |
| 322 | def forward(self, x, H, W): |
| 323 | """Forward function. |
| 324 | Args: |
| 325 | x: Input feature, tensor size (B, H*W, C). |
| 326 | H, W: Spatial resolution of the input feature. |
| 327 | """ |
| 328 | B, L, C = x.shape |
| 329 | assert L == H * W, "input feature has wrong size" |
| 330 | |
| 331 | x = x.view(B, H, W, C) |
| 332 | |
| 333 | # padding |
| 334 | pad_input = (H % 2 == 1) or (W % 2 == 1) |
| 335 | if pad_input: |
| 336 | x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) |
| 337 | |
| 338 | x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C |
| 339 | x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C |
| 340 | x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C |
| 341 | x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C |
| 342 | x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C |
| 343 | x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C |
| 344 | |
| 345 | x = self.norm(x) |
| 346 | x = self.reduction(x) |
| 347 | |
| 348 | return x |
| 349 | |
| 350 | |
| 351 | class BasicLayer(nn.Module): |