| 23 | |
| 24 | |
| 25 | class PositionalEncoding(nn.Module): |
| 26 | |
| 27 | def __init__(self, num_funcs=6): |
| 28 | super().__init__() |
| 29 | self.num_funcs = num_funcs |
| 30 | self.register_buffer('freq_bands', 2.0 ** torch.linspace(0.0, num_funcs-1, num_funcs)) |
| 31 | |
| 32 | def get_out_dim(self, in_dim): |
| 33 | return in_dim * (2 * self.num_funcs + 1) |
| 34 | |
| 35 | def forward(self, x): |
| 36 | """ |
| 37 | Args: |
| 38 | x: (..., d). |
| 39 | """ |
| 40 | shape = list(x.shape[:-1]) + [-1] |
| 41 | x = x.unsqueeze(-1) # (..., d, 1) |
| 42 | code = torch.cat([x, torch.sin(x * self.freq_bands), torch.cos(x * self.freq_bands)], dim=-1) # (..., d, 2f+1) |
| 43 | code = code.reshape(shape) |
| 44 | return code |
| 45 | |
| 46 | |
| 47 | |