2D Image to Patch Embedding
| 13 | |
| 14 | |
| 15 | class PatchEmbed(nn.Module): |
| 16 | """ 2D Image to Patch Embedding |
| 17 | """ |
| 18 | def __init__( |
| 19 | self, |
| 20 | img_size=224, |
| 21 | patch_size=16, |
| 22 | in_chans=3, |
| 23 | embed_dim=768, |
| 24 | norm_layer=None, |
| 25 | flatten=True, |
| 26 | bias=True, |
| 27 | ): |
| 28 | super().__init__() |
| 29 | img_size = to_2tuple(img_size) |
| 30 | patch_size = to_2tuple(patch_size) |
| 31 | self.img_size = img_size |
| 32 | self.patch_size = patch_size |
| 33 | self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) |
| 34 | self.num_patches = self.grid_size[0] * self.grid_size[1] |
| 35 | self.flatten = flatten |
| 36 | |
| 37 | self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias) |
| 38 | self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() |
| 39 | |
| 40 | def forward(self, x): |
| 41 | B, C, H, W = x.shape |
| 42 | _assert(H == self.img_size[0], f"Input image height ({H}) doesn't match model ({self.img_size[0]}).") |
| 43 | _assert(W == self.img_size[1], f"Input image width ({W}) doesn't match model ({self.img_size[1]}).") |
| 44 | x = self.proj(x) |
| 45 | if self.flatten: |
| 46 | x = x.flatten(2).transpose(1, 2) # BCHW -> BNC |
| 47 | x = self.norm(x) |
| 48 | return x |