An implementation of Gaussian Fourier feature mapping(random positional encoding embedding). "Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains": https://arxiv.org/abs/2006.10739 https://people.eecs.berkeley.edu/~bmild/fourfeat/index.html
| 50 | return embed, embedder_obj.out_dim |
| 51 | |
| 52 | class FourierFeatureTransform(torch.nn.Module): |
| 53 | """ |
| 54 | An implementation of Gaussian Fourier feature mapping(random positional encoding embedding). |
| 55 | "Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains": |
| 56 | https://arxiv.org/abs/2006.10739 |
| 57 | https://people.eecs.berkeley.edu/~bmild/fourfeat/index.html |
| 58 | Given an input of size [batches, num_input_channels, width, height], |
| 59 | returns a tensor of size [batches, mapping_size*2, width, height]. |
| 60 | """ |
| 61 | |
| 62 | def __init__(self, num_input_channels=3, mapping_size=256, scale=10): |
| 63 | super().__init__() |
| 64 | |
| 65 | self._num_input_channels = num_input_channels |
| 66 | self._mapping_size = mapping_size |
| 67 | B = torch.randn((num_input_channels, mapping_size)) * scale |
| 68 | B_sort = sorted(B, key=lambda x: torch.norm(x, p=2)) |
| 69 | self._B = torch.stack(B_sort) # for sape |
| 70 | |
| 71 | def forward(self, x): |
| 72 | # assert x.dim() == 4, 'Expected 4D input (got {}D input)'.format(x.dim()) |
| 73 | |
| 74 | batches, channels = x.shape |
| 75 | |
| 76 | assert channels == self._num_input_channels, \ |
| 77 | "Expected input to have {} channels (got {} channels)".format(self._num_input_channels, channels) |
| 78 | |
| 79 | # Make shape compatible for matmul with _B. |
| 80 | # From [B, C, W, H] to [(B*W*H), C]. |
| 81 | # x = x.permute(0, 2, 3, 1).reshape(batches * width * height, channels) |
| 82 | |
| 83 | res = x @ self._B.to(x.device) |
| 84 | |
| 85 | # From [(B*W*H), C] to [B, W, H, C] |
| 86 | # x = x.view(batches, width, height, self._mapping_size) |
| 87 | # From [B, W, H, C] to [B, C, W, H] |
| 88 | # x = x.permute(0, 3, 1, 2) |
| 89 | |
| 90 | res = 2 * np.pi * res |
| 91 | return torch.cat([x,torch.sin(res), torch.cos(res)], dim=1) |