89 lines
3.5 KiB
Python
89 lines
3.5 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
class FreqEncoder_torch(nn.Module):
|
|
def __init__(self, input_dim, max_freq_log2, N_freqs,
|
|
log_sampling=True, include_input=True,
|
|
periodic_fns=(torch.sin, torch.cos)):
|
|
|
|
super().__init__()
|
|
|
|
self.input_dim = input_dim
|
|
self.include_input = include_input
|
|
self.periodic_fns = periodic_fns
|
|
self.N_freqs = N_freqs
|
|
|
|
self.output_dim = 0
|
|
if self.include_input:
|
|
self.output_dim += self.input_dim
|
|
|
|
self.output_dim += self.input_dim * N_freqs * len(self.periodic_fns)
|
|
|
|
if log_sampling:
|
|
self.freq_bands = 2 ** torch.linspace(0, max_freq_log2, N_freqs)
|
|
else:
|
|
self.freq_bands = torch.linspace(2 ** 0, 2 ** max_freq_log2, N_freqs)
|
|
|
|
self.freq_bands = self.freq_bands.numpy().tolist()
|
|
|
|
def forward(self, input, max_level=None, **kwargs):
|
|
|
|
if max_level is None:
|
|
max_level = self.N_freqs
|
|
else:
|
|
max_level = int(max_level * self.N_freqs)
|
|
|
|
out = []
|
|
if self.include_input:
|
|
out.append(input)
|
|
|
|
for i in range(max_level):
|
|
freq = self.freq_bands[i]
|
|
for p_fn in self.periodic_fns:
|
|
out.append(p_fn(input * freq))
|
|
|
|
# append 0
|
|
if self.N_freqs - max_level > 0:
|
|
out.append(torch.zeros(input.shape[0], (self.N_freqs - max_level) * 2 * input.shape[1], device=input.device, dtype=input.dtype))
|
|
|
|
out = torch.cat(out, dim=-1)
|
|
|
|
return out
|
|
|
|
def get_encoder(encoding, input_dim=3,
|
|
multires=6,
|
|
degree=4,
|
|
num_levels=16, level_dim=2, base_resolution=16, log2_hashmap_size=19, desired_resolution=2048, align_corners=False, interpolation='linear',
|
|
**kwargs):
|
|
|
|
if encoding == 'None':
|
|
return lambda x, **kwargs: x, input_dim
|
|
|
|
elif encoding == 'frequency_torch':
|
|
encoder = FreqEncoder_torch(input_dim=input_dim, max_freq_log2=multires-1, N_freqs=multires, log_sampling=True)
|
|
|
|
elif encoding == 'frequency': # CUDA implementation, faster than torch.
|
|
from freqencoder import FreqEncoder
|
|
encoder = FreqEncoder(input_dim=input_dim, degree=multires)
|
|
|
|
elif encoding == 'sphere_harmonics':
|
|
from shencoder import SHEncoder
|
|
encoder = SHEncoder(input_dim=input_dim, degree=degree)
|
|
|
|
elif encoding == 'hashgrid':
|
|
from gridencoder import GridEncoder
|
|
encoder = GridEncoder(input_dim=input_dim, num_levels=num_levels, level_dim=level_dim, base_resolution=base_resolution, log2_hashmap_size=log2_hashmap_size, desired_resolution=desired_resolution, gridtype='hash', align_corners=align_corners, interpolation=interpolation)
|
|
|
|
elif encoding == 'tiledgrid':
|
|
from gridencoder import GridEncoder
|
|
encoder = GridEncoder(input_dim=input_dim, num_levels=num_levels, level_dim=level_dim, base_resolution=base_resolution, log2_hashmap_size=log2_hashmap_size, desired_resolution=desired_resolution, gridtype='tiled', align_corners=align_corners, interpolation=interpolation)
|
|
|
|
elif encoding == 'hashgrid_taichi':
|
|
from taichi_modules.hash_encoder import HashEncoderTaichi
|
|
encoder = HashEncoderTaichi(batch_size=4096) #TODO: hard encoded batch size
|
|
|
|
else:
|
|
raise NotImplementedError('Unknown encoding mode, choose from [None, frequency, sphere_harmonics, hashgrid, tiledgrid]')
|
|
|
|
return encoder, encoder.output_dim |