| 416 | |
| 417 | |
| 418 | class Direction(nn.Module): |
| 419 | def __init__(self, motion_dim): |
| 420 | super(Direction, self).__init__() |
| 421 | |
| 422 | self.weight = nn.Parameter(torch.randn(512, motion_dim)) |
| 423 | |
| 424 | def forward(self, input): |
| 425 | weight = self.weight + 1e-8 |
| 426 | Q, R = torch.linalg.qr(weight) # get eignvector, orthogonal [n1, n2, n3, n4] |
| 427 | |
| 428 | if input is None: |
| 429 | return Q |
| 430 | else: |
| 431 | input_diag = torch.diag_embed(input) # alpha, diagonal matrix |
| 432 | out = torch.matmul(input_diag, Q.T) |
| 433 | out = torch.sum(out, dim=1) |
| 434 | return out |
| 435 | |
| 436 | |
| 437 | class Synthesis(nn.Module): |