Image to Patch Embedding
| 79 | |
| 80 | |
| 81 | class PatchEmbed(nn.Module): |
| 82 | """ Image to Patch Embedding |
| 83 | """ |
| 84 | def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): |
| 85 | super().__init__() |
| 86 | img_size = to_2tuple(img_size) |
| 87 | patch_size = to_2tuple(patch_size) |
| 88 | num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) |
| 89 | self.img_size = img_size |
| 90 | self.patch_size = patch_size |
| 91 | self.num_patches = num_patches |
| 92 | |
| 93 | self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) |
| 94 | |
| 95 | def forward(self, x): |
| 96 | B, C, H, W = x.shape |
| 97 | # FIXME look at relaxing size constraints |
| 98 | # assert H == self.img_size[0] and W == self.img_size[1], \ |
| 99 | # f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." |
| 100 | x = self.proj(x).flatten(2).transpose(1, 2) |
| 101 | return x |
| 102 | |
| 103 | |
| 104 | class HybridEmbed(nn.Module): |