| 63 | |
| 64 | class PatchEmbed(nn.Module): |
| 65 | def __init__(self, in_chans, embed_dim, resolution, activation): |
| 66 | super().__init__() |
| 67 | img_size: Tuple[int, int] = to_2tuple(resolution) |
| 68 | self.patches_resolution = (img_size[0] // 4, img_size[1] // 4) |
| 69 | self.num_patches = self.patches_resolution[0] * \ |
| 70 | self.patches_resolution[1] |
| 71 | self.in_chans = in_chans |
| 72 | self.embed_dim = embed_dim |
| 73 | n = embed_dim |
| 74 | self.seq = nn.Sequential( |
| 75 | Conv2d_BN(in_chans, n // 2, 3, 2, 1), |
| 76 | activation(), |
| 77 | Conv2d_BN(n // 2, n, 3, 2, 1), |
| 78 | ) |
| 79 | |
| 80 | def forward(self, x): |
| 81 | return self.seq(x) |