| 113 | |
| 114 | class PatchEmbed(nn.Module): |
| 115 | def __init__(self, img_size=108, patch_size=9, in_channels=3, embed_dim=768): |
| 116 | super().__init__() |
| 117 | img_size = to_2tuple(img_size) |
| 118 | patch_size = to_2tuple(patch_size) |
| 119 | num_patches = (img_size[1] // patch_size[1]) * \ |
| 120 | (img_size[0] // patch_size[0]) |
| 121 | self.img_size = img_size |
| 122 | self.patch_size = patch_size |
| 123 | self.num_patches = num_patches |
| 124 | self.proj = nn.Conv2d(in_channels, embed_dim, |
| 125 | kernel_size=patch_size, stride=patch_size) |
| 126 | |
| 127 | def forward(self, x): |
| 128 | batch_size, channels, height, width = x.shape |