| 18 | |
| 19 | # Positional encoding (section 5.1) |
| 20 | class Embedder: |
| 21 | def __init__(self, **kwargs): |
| 22 | self.kwargs = kwargs |
| 23 | self.create_embedding_fn() |
| 24 | |
| 25 | def create_embedding_fn(self): |
| 26 | embed_fns = [] |
| 27 | d = self.kwargs['input_dims'] |
| 28 | out_dim = 0 |
| 29 | if self.kwargs['include_input']: |
| 30 | embed_fns.append(lambda x : x) |
| 31 | out_dim += d |
| 32 | |
| 33 | max_freq = self.kwargs['max_freq_log2'] |
| 34 | N_freqs = self.kwargs['num_freqs'] |
| 35 | |
| 36 | if self.kwargs['log_sampling']: |
| 37 | freq_bands = 2.**torch.linspace(0., max_freq, steps=N_freqs) |
| 38 | else: |
| 39 | freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=N_freqs) |
| 40 | |
| 41 | for freq in freq_bands: |
| 42 | for p_fn in self.kwargs['periodic_fns']: |
| 43 | embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq)) |
| 44 | out_dim += d |
| 45 | |
| 46 | self.embed_fns = embed_fns |
| 47 | self.out_dim = out_dim |
| 48 | |
| 49 | def embed(self, inputs): |
| 50 | return torch.cat([fn(inputs) for fn in self.embed_fns], -1) |
| 51 | |
| 52 | |
| 53 | def get_embedder(multires, i=0): |