Spaces:
Running on Zero
Running on Zero
File size: 2,507 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | import torch
from einops import rearrange, repeat
# TODO: Try RoPE / AliBi
def sincos_1d(L, D, device, dtype):
# 1d sinusoidal position encoding where element j of ith patch embedding is encoded as:
# PE[i, 2j] = sin(i / 10000^(2j/D)) # even indices
# PE[i, 2j+1] = cos(i / 10000^(2j/D)) # odd indices
assert D % 2 == 0, "Encoding dimension must be even"
# position indices [L, 1] and dimension indices [1, D/2]
pos = rearrange(torch.arange(L, device=device, dtype=dtype), 'l -> l 1') # [L,1]
i = rearrange(torch.arange(D // 2, device=device, dtype=dtype), 'd -> 1 d')# [1,D/2]
# angular frequencies: 1/10000^(2i/D) for each dimension
div = torch.pow(torch.tensor(10000.0, device=device, dtype=dtype), (2*i)/D)
# angles: pos * freq for each position-dimension pair
angles = pos / div # [L, D/2] (broadcasted together)
pe = torch.zeros(L, D, device=device, dtype=dtype)
pe[:, 0::2] = torch.sin(angles) # even indices
pe[:, 1::2] = torch.cos(angles) # odd indices
return pe # [L, D]
def sincos_time(T, D, device, dtype):
# temporal PE (1d sinusoidal PE across time)
return sincos_1d(T, D, device, dtype) # reuse the same 1D builder
def build_spatial_only_pe(frame_size, patch_size, embed_dim, device='cpu', dtype=torch.float32):
# spatial positional encodings for a grid of patches in first 2/3 of embed dim (evenly into x and y axes)
# last 1/3 for temporal PE padded with 0s
H, W = frame_size
Hp, Wp = H // patch_size, W // patch_size
# split dimensions (ensure temporal even)
temporal_dim = (embed_dim // 3) & ~1
spatial_dims = embed_dim - temporal_dim
# split spatial dims between x and y (ensure both even)
spatial_x_dim = (spatial_dims // 2) & ~1
spatial_y_dim = spatial_dims - spatial_x_dim
assert spatial_x_dim % 2 == 0 and spatial_y_dim % 2 == 0 and temporal_dim % 2 == 0
# 2d PE for x and y axes
pe_x = sincos_1d(Wp, spatial_x_dim, device, dtype) # [Wp, Dx]
pe_y = sincos_1d(Hp, spatial_y_dim, device, dtype) # [Hp, Dy]
pe_x = repeat(pe_x, 'wp dx -> hp wp dx', hp=Hp) # [Hp, Wp, Dx]
pe_y = repeat(pe_y, 'hp dy -> hp wp dy', wp=Wp) # [Hp, Wp, Dy]
pe_spatial = torch.cat([
pe_x,
pe_y,
torch.zeros(Hp, Wp, temporal_dim, device=device, dtype=dtype) # zero temporal tail
], dim=-1) # [Hp, Wp, E]
pe_spatial = rearrange(pe_spatial, 'hp wp e -> 1 (hp wp) e') # [1, P, E]
return pe_spatial # [1, P, E] |