| 29 | |
| 30 | |
| 31 | def build_spatial_only_pe(frame_size, patch_size, embed_dim, device='cpu', dtype=torch.float32): |
| 32 | # spatial positional encodings for a grid of patches in first 2/3 of embed dim (evenly into x and y axes) |
| 33 | # last 1/3 for temporal PE padded with 0s |
| 34 | H, W = frame_size |
| 35 | Hp, Wp = H // patch_size, W // patch_size |
| 36 | |
| 37 | # split dimensions (ensure temporal even) |
| 38 | temporal_dim = (embed_dim // 3) & ~1 |
| 39 | spatial_dims = embed_dim - temporal_dim |
| 40 | |
| 41 | # split spatial dims between x and y (ensure both even) |
| 42 | spatial_x_dim = (spatial_dims // 2) & ~1 |
| 43 | spatial_y_dim = spatial_dims - spatial_x_dim |
| 44 | |
| 45 | assert spatial_x_dim % 2 == 0 and spatial_y_dim % 2 == 0 and temporal_dim % 2 == 0 |
| 46 | |
| 47 | # 2d PE for x and y axes |
| 48 | pe_x = sincos_1d(Wp, spatial_x_dim, device, dtype) # [Wp, Dx] |
| 49 | pe_y = sincos_1d(Hp, spatial_y_dim, device, dtype) # [Hp, Dy] |
| 50 | pe_x = repeat(pe_x, 'wp dx -> hp wp dx', hp=Hp) # [Hp, Wp, Dx] |
| 51 | pe_y = repeat(pe_y, 'hp dy -> hp wp dy', wp=Wp) # [Hp, Wp, Dy] |
| 52 | |
| 53 | pe_spatial = torch.cat([ |
| 54 | pe_x, |
| 55 | pe_y, |
| 56 | torch.zeros(Hp, Wp, temporal_dim, device=device, dtype=dtype) # zero temporal tail |
| 57 | ], dim=-1) # [Hp, Wp, E] |
| 58 | |
| 59 | pe_spatial = rearrange(pe_spatial, 'hp wp e -> 1 (hp wp) e') # [1, P, E] |
| 60 | return pe_spatial # [1, P, E] |