Forward function. Args: x: Input feature, tensor size (B, H*W, C). H, W: Spatial resolution of the input feature.
(self, x, H, W)
| 262 | self.downsample = None |
| 263 | |
| 264 | def forward(self, x, H, W): |
| 265 | """ Forward function. |
| 266 | |
| 267 | Args: |
| 268 | x: Input feature, tensor size (B, H*W, C). |
| 269 | H, W: Spatial resolution of the input feature. |
| 270 | """ |
| 271 | for blk in self.blocks: |
| 272 | blk.H, blk.W = H, W |
| 273 | if self.use_checkpoint: |
| 274 | x = checkpoint.checkpoint(blk, x) |
| 275 | else: |
| 276 | x = blk(x) |
| 277 | if self.downsample is not None: |
| 278 | x_reshaped = x.transpose(1, 2).view(x.shape[0], x.shape[-1], H, W) |
| 279 | x_down = self.downsample(x_reshaped) |
| 280 | x_down = x_down.flatten(2).transpose(1, 2) |
| 281 | Wh, Ww = (H + 1) // 2, (W + 1) // 2 |
| 282 | return x, H, W, x_down, Wh, Ww |
| 283 | else: |
| 284 | return x, H, W, x, H, W |
| 285 | |
| 286 | |
| 287 | class PatchEmbed(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected