Spaces:
Sleeping
Sleeping
File size: 1,566 Bytes
78627b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | import torch
import torch.nn as nn
from typing import Tuple
class PatchEmbedding(nn.Module):
"""
Convert image to patch embeddings.
Args:
img_size (int): Size of input image (assumed square).
patch_size (int): Size of each patch (assumed square).
in_chans (int): Number of input channels.
embed_dim (int): Embedding dimension.
"""
def __init__(
self,
img_size: int = 224,
patch_size: int = 16,
in_chans: int = 3,
embed_dim: int = 768
):
super().__init__()
self.img_size = img_size
self.patch_size = patch_size
self.grid_size = img_size // patch_size
self.num_patches = self.grid_size ** 2
self.proj = nn.Conv2d(
in_chans,
embed_dim,
kernel_size=patch_size,
stride=patch_size
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x (torch.Tensor): Input image tensor of shape (B, C, H, W).
Returns:
torch.Tensor: Patch embeddings of shape (B, L, D) where L is num_patches.
"""
B, C, H, W = x.shape
assert H == self.img_size and W == self.img_size, \
f"Input image size ({H}x{W}) doesn't match model image size ({self.img_size}x{self.img_size})"
# (B, C, H, W) -> (B, D, H/P, W/P)
x = self.proj(x)
# (B, D, H/P, W/P) -> (B, D, L) -> (B, L, D)
x = x.flatten(2).transpose(1, 2)
return x
|