| 581 | |
| 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. |
| 618 | |
| 619 | Args: |
| 620 | tensor (torch.Tensor): Input tensor to be positionally encoded. |
| 621 | encoding_size (optional, int): Number of encoding functions used to compute |
| 622 | a positional encoding (default: 6). |
| 623 | include_input (optional, bool): Whether or not to include the input in the |
| 624 | positional encoding (default: True). |
| 625 | |
| 626 | Returns: |
| 627 | (torch.Tensor): Positional encoding of the input tensor. |
| 628 | """ |
| 629 | |
| 630 | encoding = [tensor] if self.include_input else [] |
| 631 | if self.gaussian_pe: |
| 632 | for func in [torch.sin, torch.cos]: |
| 633 | encoding.append(func(torch.matmul(tensor, self.gaussian_weights.T))) |
| 634 | else: |
| 635 | for idx, freq in enumerate(self.frequency_bands): |
| 636 | for func in [torch.sin, torch.cos]: |
| 637 | if self.normalization is not None: |
| 638 | encoding.append(self.normalization[idx]*func(tensor * freq)) |
| 639 | else: |
| 640 | encoding.append(func(tensor * freq)) |