| 229 | |
| 230 | @persistence.persistent_class |
| 231 | class TemporalInput(nn.Module): |
| 232 | def __init__(self, cfg: DictConfig, channel_dim: int, motion_v_dim: int): |
| 233 | super().__init__() |
| 234 | |
| 235 | self.cfg = cfg |
| 236 | self.motion_v_dim = motion_v_dim |
| 237 | self.const = nn.Parameter(torch.randn(1, channel_dim, 4, 4)) |
| 238 | |
| 239 | def get_dim(self): |
| 240 | return self.motion_v_dim + self.const.shape[1] |
| 241 | |
| 242 | def forward(self, motion_v: torch.Tensor) -> torch.Tensor: |
| 243 | """ |
| 244 | motion_v: [batch_size, motion_v_dim] |
| 245 | """ |
| 246 | out = torch.cat([ |
| 247 | self.const.repeat(len(motion_v), 1, 1, 1), |
| 248 | motion_v.unsqueeze(2).unsqueeze(3).repeat(1, 1, *self.const.shape[2:]), |
| 249 | ], dim=1) # [batch_size, channel_dim + num_fourier_feats * 2] |
| 250 | |
| 251 | return out |
| 252 | |
| 253 | #---------------------------------------------------------------------------- |
| 254 | |