| 4 | from models.positional_encoding import build_spatial_only_pe |
| 5 | |
| 6 | class PatchEmbedding(nn.Module): |
| 7 | def __init__(self, frame_size=(128, 128), patch_size=8, embed_dim=128): |
| 8 | super().__init__() |
| 9 | H, W = frame_size |
| 10 | self.frame_size = frame_size |
| 11 | self.patch_size = patch_size |
| 12 | self.embed_dim = embed_dim |
| 13 | self.Hp, self.Wp = H // patch_size, W // patch_size |
| 14 | self.num_patches = self.Hp * self.Wp |
| 15 | |
| 16 | # split embed dim into thirds for spatial x, spatial y, and temporal |
| 17 | base_split = (embed_dim // 3) & ~1 |
| 18 | remaining_dim = embed_dim - base_split |
| 19 | self.spatial_x_dim = (remaining_dim // 2) & ~1 |
| 20 | self.spatial_y_dim = remaining_dim - self.spatial_x_dim |
| 21 | self.temporal_dim = base_split |
| 22 | |
| 23 | # ensure the embed dim is split wholy into thirds and each third is even |
| 24 | assert (self.spatial_x_dim + self.spatial_y_dim + self.temporal_dim) == embed_dim, \ |
| 25 | f"Dimension mismatch: {self.spatial_x_dim} + {self.spatial_y_dim} + {self.temporal_dim} != {embed_dim}" |
| 26 | assert self.spatial_x_dim % 2 == 0 and self.spatial_y_dim % 2 == 0 and self.temporal_dim % 2 == 0, \ |
| 27 | f"Embed dim x={self.spatial_x_dim}, y={self.spatial_y_dim}, t={self.temporal_dim}" |
| 28 | |
| 29 | pe_spatial = build_spatial_only_pe(self.frame_size, self.patch_size, self.embed_dim, device='cpu', dtype=torch.float32) # [1,P,E] |
| 30 | self.register_buffer("pos_spatial", pe_spatial, persistent=False) |
| 31 | |
| 32 | # pixel patches to embeddings |
| 33 | self.proj = nn.Conv2d(3 * self.patch_size * self.patch_size, self.embed_dim, 1) |
| 34 | |
| 35 | |
| 36 | def forward(self, frames): |
| 37 | B, T, C, H, W = frames.shape |
| 38 | # go from frames to patches |
| 39 | x = rearrange(frames, 'b t c (hp p1) (wp p2) -> (b t) (c p1 p2) hp wp', p1=self.patch_size, p2=self.patch_size) # [(B*T), 3*p*p, Hp, Wp] |
| 40 | x = self.proj(x) # [(B*T), E, Hp, Wp] |
| 41 | x = rearrange(x, '(b t) e hp wp -> b t (hp wp) e', b=B, t=T) # [B, T, P, E] |
| 42 | # add 2d spatial pos encoding (first 2/3 of embed dim) |
| 43 | x = x + self.pos_spatial.to(dtype=x.dtype, device=x.device) # [B, T, P, E] |
| 44 | return x |