| import torch |
| import torch.nn as nn |
| from timm.layers import PatchEmbed as TimmPatchEmbed |
|
|
|
|
| class PatchEmbed(nn.Module): |
| """ |
| 2D Image to Patch Embedding. |
| |
| Args: |
| img_size (tuple[int, int]): Input image size (H, W). |
| patch_size (tuple[int, int]): Patch size (H, W). |
| in_chans (int): Number of input channels. |
| embed_dim (int): Embedding dimension. |
| """ |
|
|
| def __init__( |
| self, |
| img_size: tuple[int, int] = (128, 256), |
| patch_size: tuple[int, int] = (16, 16), |
| in_chans: int = 1, |
| embed_dim: int = 768, |
| ): |
| super().__init__() |
| self.img_size = img_size |
| self.patch_size = patch_size |
| self.in_chans = in_chans |
| self.embed_dim = embed_dim |
|
|
| self.patch_embed = TimmPatchEmbed( |
| img_size=img_size, |
| patch_size=patch_size, |
| in_chans=in_chans, |
| embed_dim=embed_dim, |
| flatten=True, |
| bias=True, |
| strict_img_size=False, |
| ) |
| self.num_patches = self.patch_embed.num_patches |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """ |
| Forward pass. |
| |
| Args: |
| x (torch.Tensor): Input tensor [B, C, H, W]. |
| |
| Returns: |
| torch.Tensor: Patch embeddings [B, N, D]. |
| """ |
| return self.patch_embed(x) |
|
|
