Image to Patch Embedding
| 242 | |
| 243 | |
| 244 | class PatchEmbed(nn.Module): |
| 245 | """ Image to Patch Embedding |
| 246 | """ |
| 247 | def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): |
| 248 | super().__init__() |
| 249 | num_patches = (img_size // patch_size) * (img_size // patch_size) |
| 250 | self.img_size = img_size |
| 251 | self.patch_size = patch_size |
| 252 | self.num_patches = num_patches |
| 253 | |
| 254 | self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) |
| 255 | |
| 256 | def forward(self, x): |
| 257 | B, C, H, W = x.shape |
| 258 | x = self.proj(x).flatten(2).transpose(1, 2) |
| 259 | return x |
| 260 | |
| 261 | |
| 262 | class VisionTransformer(nn.Module): |