| 300 | |
| 301 | @persistence.persistent_class |
| 302 | class FixedTimeEncoder(nn.Module): |
| 303 | def __init__(self, |
| 304 | max_num_frames: int, # Maximum T size |
| 305 | skip_small_t_freqs: int=0, # How many high frequencies we should skip |
| 306 | ): |
| 307 | super().__init__() |
| 308 | |
| 309 | assert max_num_frames >= 1, f"Wrong max_num_frames: {max_num_frames}" |
| 310 | fourier_coefs = construct_log_spaced_freqs(max_num_frames, skip_small_t_freqs=skip_small_t_freqs) |
| 311 | self.register_buffer('fourier_coefs', fourier_coefs) # [1, num_fourier_feats] |
| 312 | |
| 313 | def get_dim(self) -> int: |
| 314 | return self.fourier_coefs.shape[1] * 2 |
| 315 | |
| 316 | def forward(self, t: torch.Tensor) -> torch.Tensor: |
| 317 | assert t.ndim == 2, f"Wrong shape: {t.shape}" |
| 318 | |
| 319 | t = t.view(-1).float() # [batch_size * num_frames] |
| 320 | fourier_raw_embs = self.fourier_coefs * t.unsqueeze(1) # [bf, num_fourier_feats] |
| 321 | |
| 322 | fourier_embs = torch.cat([ |
| 323 | fourier_raw_embs.sin(), |
| 324 | fourier_raw_embs.cos(), |
| 325 | ], dim=1) # [bf, num_fourier_feats * 2] |
| 326 | |
| 327 | return fourier_embs |
| 328 | |
| 329 | #---------------------------------------------------------------------------- |
| 330 | |