| 66 | self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() |
| 67 | |
| 68 | def forward(self, x: Tensor) -> Tensor: |
| 69 | _, _, H, W = x.shape |
| 70 | patch_H, patch_W = self.patch_size |
| 71 | |
| 72 | assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}" |
| 73 | assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}" |
| 74 | |
| 75 | x = self.proj(x) # B C H W |
| 76 | H, W = x.size(2), x.size(3) |
| 77 | x = x.flatten(2).transpose(1, 2) # B HW C |
| 78 | x = self.norm(x) |
| 79 | if not self.flatten_embedding: |
| 80 | x = x.reshape(-1, H, W, self.embed_dim) # B H W C |
| 81 | return x |
| 82 | |
| 83 | def flops(self) -> float: |
| 84 | Ho, Wo = self.patches_resolution |