Forward function. Args: x: Input feature, tensor size (B, H*W, C). H, W: Spatial resolution of the input feature.
(self, x, H, W)
| 309 | self.norm = norm_layer(4 * dim) |
| 310 | |
| 311 | def forward(self, x, H, W): |
| 312 | """Forward function. |
| 313 | Args: |
| 314 | x: Input feature, tensor size (B, H*W, C). |
| 315 | H, W: Spatial resolution of the input feature. |
| 316 | """ |
| 317 | B, L, C = x.shape |
| 318 | assert L == H * W, "input feature has wrong size" |
| 319 | |
| 320 | x = x.view(B, H, W, C) |
| 321 | |
| 322 | # padding |
| 323 | pad_input = (H % 2 == 1) or (W % 2 == 1) |
| 324 | if pad_input: |
| 325 | x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) |
| 326 | |
| 327 | x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C |
| 328 | x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C |
| 329 | x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C |
| 330 | x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C |
| 331 | x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C |
| 332 | x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C |
| 333 | |
| 334 | x = self.norm(x) |
| 335 | x = self.reduction(x) |
| 336 | |
| 337 | return x |
| 338 | |
| 339 | |
| 340 | class BasicLayer(nn.Module): |