| 582 | |
| 583 | class PositionalEncoding(nn.Module): |
| 584 | def __init__(self, num_encoding_functions=6, include_input=True, log_sampling=True, normalize=False, |
| 585 | input_dim=3, gaussian_pe=False, gaussian_variance=38): |
| 586 | super().__init__() |
| 587 | self.num_encoding_functions = num_encoding_functions |
| 588 | self.include_input = include_input |
| 589 | self.log_sampling = log_sampling |
| 590 | self.normalize = normalize |
| 591 | self.gaussian_pe = gaussian_pe |
| 592 | self.normalization = None |
| 593 | |
| 594 | if self.gaussian_pe: |
| 595 | # this needs to be registered as a parameter so that it is saved in the model state dict |
| 596 | # and so that it is converted using .cuda(). Doesn't need to be trained though |
| 597 | self.gaussian_weights = nn.Parameter(gaussian_variance * torch.randn(num_encoding_functions, input_dim), |
| 598 | requires_grad=False) |
| 599 | |
| 600 | else: |
| 601 | self.frequency_bands = None |
| 602 | if self.log_sampling: |
| 603 | self.frequency_bands = 2.0 ** torch.linspace( |
| 604 | 0.0, |
| 605 | self.num_encoding_functions - 1, |
| 606 | self.num_encoding_functions) |
| 607 | else: |
| 608 | self.frequency_bands = torch.linspace( |
| 609 | 2.0 ** 0.0, |
| 610 | 2.0 ** (self.num_encoding_functions - 1), |
| 611 | self.num_encoding_functions) |
| 612 | |
| 613 | if normalize: |
| 614 | self.normalization = torch.tensor(1/self.frequency_bands) |
| 615 | |
| 616 | def forward(self, tensor) -> torch.Tensor: |
| 617 | r"""Apply positional encoding to the input. |