| 22 | from timm.models.layers import Mlp, DropPath, activations, to_2tuple, trunc_normal_ |
| 23 | |
| 24 | class PatchEmbed(nn.Module): |
| 25 | def __init__(self, patch_size=16, in_chans=1, embed_dim=64, pe='linear'): |
| 26 | super().__init__() |
| 27 | self.patch_size = patch_size |
| 28 | self.dim = embed_dim |
| 29 | self.pe = pe |
| 30 | |
| 31 | # assert patch_size == 8 |
| 32 | if patch_size == 8: |
| 33 | self.proj = nn.Sequential( |
| 34 | nn.Conv2d(in_chans, embed_dim//4, kernel_size=6, stride=2, padding=2), |
| 35 | nn.ReLU(), |
| 36 | nn.Conv2d(embed_dim//4, embed_dim//2, kernel_size=6, stride=2, padding=2), |
| 37 | nn.ReLU(), |
| 38 | nn.Conv2d(embed_dim//2, embed_dim, kernel_size=6, stride=2, padding=2), |
| 39 | ) |
| 40 | elif patch_size == 4: |
| 41 | self.proj = nn.Sequential( |
| 42 | nn.Conv2d(in_chans, embed_dim//4, kernel_size=6, stride=2, padding=2), |
| 43 | nn.ReLU(), |
| 44 | nn.Conv2d(embed_dim//4, embed_dim, kernel_size=6, stride=2, padding=2), |
| 45 | ) |
| 46 | else: |
| 47 | print(f"patch size = {patch_size} is unacceptable.") |
| 48 | |
| 49 | self.ffn_with_coord = nn.Sequential( |
| 50 | nn.Conv2d(embed_dim*2, embed_dim*2, kernel_size=1), |
| 51 | nn.ReLU(), |
| 52 | nn.Conv2d(embed_dim*2, embed_dim*2, kernel_size=1) |
| 53 | ) |
| 54 | self.norm = nn.LayerNorm(embed_dim*2) |
| 55 | |
| 56 | def forward(self, x) -> Tuple[torch.Tensor, Size_]: |
| 57 | B, C, H, W = x.shape # C == 1 |
| 58 | |
| 59 | pad_l = pad_t = 0 |
| 60 | pad_r = (self.patch_size - W % self.patch_size) % self.patch_size |
| 61 | pad_b = (self.patch_size - H % self.patch_size) % self.patch_size |
| 62 | x = F.pad(x, (pad_l, pad_r, pad_t, pad_b)) |
| 63 | |
| 64 | x = self.proj(x) |
| 65 | out_size = x.shape[2:] |
| 66 | |
| 67 | patch_coord = coords_grid(B, out_size[0], out_size[1]).to(x.device) * self.patch_size + self.patch_size/2 # in feature coordinate space |
| 68 | patch_coord = patch_coord.view(B, 2, -1).permute(0, 2, 1) |
| 69 | if self.pe == 'linear': |
| 70 | patch_coord_enc = LinearPositionEmbeddingSine(patch_coord, dim=self.dim) |
| 71 | elif self.pe == 'exp': |
| 72 | patch_coord_enc = ExpPositionEmbeddingSine(patch_coord, dim=self.dim) |
| 73 | patch_coord_enc = patch_coord_enc.permute(0, 2, 1).view(B, -1, out_size[0], out_size[1]) |
| 74 | |
| 75 | x_pe = torch.cat([x, patch_coord_enc], dim=1) |
| 76 | x = self.ffn_with_coord(x_pe) |
| 77 | x = self.norm(x.flatten(2).transpose(1, 2)) |
| 78 | |
| 79 | return x, out_size |
| 80 | |
| 81 | from .twins import Block, CrossBlock |