Image to Patch Embedding
| 226 | |
| 227 | |
| 228 | class PatchEmbed(nn.Module): |
| 229 | """ Image to Patch Embedding |
| 230 | """ |
| 231 | def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, std=False): |
| 232 | super().__init__() |
| 233 | img_size = to_2tuple(img_size) |
| 234 | patch_size = to_2tuple(patch_size) |
| 235 | num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) |
| 236 | self.img_size = img_size |
| 237 | self.patch_size = patch_size |
| 238 | self.num_patches = num_patches |
| 239 | self.norm = nn.LayerNorm(embed_dim) |
| 240 | if std: |
| 241 | self.proj = conv_3xnxn_std(in_chans, embed_dim, kernel_size=patch_size[0], stride=patch_size[0]) |
| 242 | else: |
| 243 | self.proj = conv_1xnxn(in_chans, embed_dim, kernel_size=patch_size[0], stride=patch_size[0]) |
| 244 | |
| 245 | def forward(self, x): |
| 246 | B, C, T, H, W = x.shape |
| 247 | # FIXME look at relaxing size constraints |
| 248 | # assert H == self.img_size[0] and W == self.img_size[1], \ |
| 249 | # f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." |
| 250 | x = self.proj(x) |
| 251 | B, C, T, H, W = x.shape |
| 252 | x = x.flatten(2).transpose(1, 2) |
| 253 | x = self.norm(x) |
| 254 | x = x.reshape(B, T, H, W, -1).permute(0, 4, 1, 2, 3).contiguous() |
| 255 | return x |
| 256 | |
| 257 | |
| 258 | @MODEL_REGISTRY.register() |