| 543 | |
| 544 | |
| 545 | class FFPositionalEncoding(nn.Module): |
| 546 | def __init__(self, embedding_size, scale, dims=2, gaussian=True): |
| 547 | super().__init__() |
| 548 | self.embedding_size = embedding_size |
| 549 | self.scale = scale |
| 550 | |
| 551 | if gaussian: |
| 552 | bvals = torch.randn(embedding_size // 2, dims) * scale |
| 553 | else: |
| 554 | bvals = 2.**torch.linspace(0, scale, embedding_size//2) - 1 |
| 555 | |
| 556 | if dims == 1: |
| 557 | bvals = bvals[:, None] |
| 558 | |
| 559 | elif dims == 2: |
| 560 | bvals = torch.stack([bvals, torch.zeros_like(bvals)], dim=-1) |
| 561 | bvals = torch.cat([bvals, torch.roll(bvals, 1, -1)], dim=0) |
| 562 | |
| 563 | else: |
| 564 | tmp = (dims-1)*(torch.zeros_like(bvals),) |
| 565 | bvals = torch.stack([bvals, *tmp], dim=-1) |
| 566 | |
| 567 | tmp = [torch.roll(bvals, i, -1) for i in range(1, dims)] |
| 568 | bvals = torch.cat([bvals, *tmp], dim=0) |
| 569 | |
| 570 | avals = torch.ones((bvals.shape[0])) |
| 571 | self.avals = nn.Parameter(avals, requires_grad=False) |
| 572 | self.bvals = nn.Parameter(bvals, requires_grad=False) |
| 573 | |
| 574 | def forward(self, tensor) -> torch.Tensor: |
| 575 | """ |
| 576 | Apply positional encoding to the input. |
| 577 | """ |
| 578 | |
| 579 | return torch.cat([self.avals * torch.sin((2.*np.pi*tensor) @ self.bvals.T), |
| 580 | self.avals * torch.cos((2.*np.pi*tensor) @ self.bvals.T)], dim=-1) |
| 581 | |
| 582 | |
| 583 | class PositionalEncoding(nn.Module): |