2D image to patch embedding: (B,C,H,W) -> (B,N,D) Args: img_size: Image size. patch_size: Patch token size. in_chans: Number of input image channels. embed_dim: Number of linear projection output channels. norm_layer: Normalization layer.
| 23 | |
| 24 | |
| 25 | class PatchEmbed(nn.Module): |
| 26 | """ |
| 27 | 2D image to patch embedding: (B,C,H,W) -> (B,N,D) |
| 28 | |
| 29 | Args: |
| 30 | img_size: Image size. |
| 31 | patch_size: Patch token size. |
| 32 | in_chans: Number of input image channels. |
| 33 | embed_dim: Number of linear projection output channels. |
| 34 | norm_layer: Normalization layer. |
| 35 | """ |
| 36 | |
| 37 | def __init__( |
| 38 | self, |
| 39 | img_size: Union[int, Tuple[int, int]] = 224, |
| 40 | patch_size: Union[int, Tuple[int, int]] = 16, |
| 41 | in_chans: int = 3, |
| 42 | embed_dim: int = 768, |
| 43 | norm_layer: Optional[Callable] = None, |
| 44 | flatten_embedding: bool = True, |
| 45 | ) -> None: |
| 46 | super().__init__() |
| 47 | |
| 48 | image_HW = make_2tuple(img_size) |
| 49 | patch_HW = make_2tuple(patch_size) |
| 50 | patch_grid_size = (image_HW[0] // patch_HW[0], image_HW[1] // patch_HW[1]) |
| 51 | |
| 52 | self.img_size = image_HW |
| 53 | self.patch_size = patch_HW |
| 54 | self.patches_resolution = patch_grid_size |
| 55 | self.num_patches = patch_grid_size[0] * patch_grid_size[1] |
| 56 | |
| 57 | self.in_chans = in_chans |
| 58 | self.embed_dim = embed_dim |
| 59 | |
| 60 | self.flatten_embedding = flatten_embedding |
| 61 | |
| 62 | self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW) |
| 63 | self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() |
| 64 | |
| 65 | def forward(self, x: Tensor) -> Tensor: |
| 66 | _, _, H, W = x.shape |
| 67 | patch_H, patch_W = self.patch_size |
| 68 | |
| 69 | assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}" |
| 70 | assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}" |
| 71 | |
| 72 | x = self.proj(x) # B C H W |
| 73 | H, W = x.size(2), x.size(3) |
| 74 | x = x.flatten(2).transpose(1, 2) # B HW C |
| 75 | x = self.norm(x) |
| 76 | if not self.flatten_embedding: |
| 77 | x = x.reshape(-1, H, W, self.embed_dim) # B H W C |
| 78 | return x |
| 79 | |
| 80 | def flops(self) -> float: |
| 81 | Ho, Wo = self.patches_resolution |
| 82 | flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) |