Spaces:
Running on Zero
Running on Zero
File size: 2,208 Bytes
dcc4594 | 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 | import torch
import torch.nn as nn
from einops import rearrange
from models.positional_encoding import build_spatial_only_pe
class PatchEmbedding(nn.Module):
def __init__(self, frame_size=(128, 128), patch_size=8, embed_dim=128):
super().__init__()
H, W = frame_size
self.frame_size = frame_size
self.patch_size = patch_size
self.embed_dim = embed_dim
self.Hp, self.Wp = H // patch_size, W // patch_size
self.num_patches = self.Hp * self.Wp
# split embed dim into thirds for spatial x, spatial y, and temporal
base_split = (embed_dim // 3) & ~1
remaining_dim = embed_dim - base_split
self.spatial_x_dim = (remaining_dim // 2) & ~1
self.spatial_y_dim = remaining_dim - self.spatial_x_dim
self.temporal_dim = base_split
# ensure the embed dim is split wholy into thirds and each third is even
assert (self.spatial_x_dim + self.spatial_y_dim + self.temporal_dim) == embed_dim, \
f"Dimension mismatch: {self.spatial_x_dim} + {self.spatial_y_dim} + {self.temporal_dim} != {embed_dim}"
assert self.spatial_x_dim % 2 == 0 and self.spatial_y_dim % 2 == 0 and self.temporal_dim % 2 == 0, \
f"Embed dim x={self.spatial_x_dim}, y={self.spatial_y_dim}, t={self.temporal_dim}"
pe_spatial = build_spatial_only_pe(self.frame_size, self.patch_size, self.embed_dim, device='cpu', dtype=torch.float32) # [1,P,E]
self.register_buffer("pos_spatial", pe_spatial, persistent=False)
# pixel patches to embeddings
self.proj = nn.Conv2d(3 * self.patch_size * self.patch_size, self.embed_dim, 1)
def forward(self, frames):
B, T, C, H, W = frames.shape
# go from frames to patches
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]
x = self.proj(x) # [(B*T), E, Hp, Wp]
x = rearrange(x, '(b t) e hp wp -> b t (hp wp) e', b=B, t=T) # [B, T, P, E]
# add 2d spatial pos encoding (first 2/3 of embed dim)
x = x + self.pos_spatial.to(dtype=x.dtype, device=x.device) # [B, T, P, E]
return x
|