Forward function. Args: x: Input feature, tensor size (B, H*W, C). H, W: Spatial resolution of the input feature.
(self, x, H, W)
| 273 | self.downsample = None |
| 274 | |
| 275 | def forward(self, x, H, W): |
| 276 | """ Forward function. |
| 277 | |
| 278 | Args: |
| 279 | x: Input feature, tensor size (B, H*W, C). |
| 280 | H, W: Spatial resolution of the input feature. |
| 281 | """ |
| 282 | for blk in self.blocks: |
| 283 | blk.H, blk.W = H, W |
| 284 | if self.use_checkpoint: |
| 285 | x = checkpoint.checkpoint(blk, x) |
| 286 | else: |
| 287 | x = blk(x) |
| 288 | if self.downsample is not None: |
| 289 | x_reshaped = x.transpose(1, 2).view(x.shape[0], x.shape[-1], H, W) |
| 290 | x_down = self.downsample(x_reshaped) |
| 291 | x_down = x_down.flatten(2).transpose(1, 2) |
| 292 | Wh, Ww = (H + 1) // 2, (W + 1) // 2 |
| 293 | return x, H, W, x_down, Wh, Ww |
| 294 | else: |
| 295 | return x, H, W, x, H, W |
| 296 | |
| 297 | |
| 298 | # class PatchEmbed(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected