|
from torch import nn, optim |
|
import math |
|
import torch |
|
import torch.nn.functional as F |
|
from einops import rearrange |
|
import numpy as np |
|
from datetime import datetime |
|
import positional_encoding as PE |
|
|
|
""" |
|
FCNet |
|
""" |
|
class ResLayer(nn.Module): |
|
def __init__(self, linear_size): |
|
super(ResLayer, self).__init__() |
|
self.l_size = linear_size |
|
self.nonlin1 = nn.ReLU(inplace=True) |
|
self.nonlin2 = nn.ReLU(inplace=True) |
|
self.dropout1 = nn.Dropout() |
|
self.w1 = nn.Linear(self.l_size, self.l_size) |
|
self.w2 = nn.Linear(self.l_size, self.l_size) |
|
|
|
def forward(self, x): |
|
y = self.w1(x) |
|
y = self.nonlin1(y) |
|
y = self.dropout1(y) |
|
y = self.w2(y) |
|
y = self.nonlin2(y) |
|
out = x + y |
|
|
|
return out |
|
|
|
class FCNet(nn.Module): |
|
def __init__(self, num_inputs, num_classes, dim_hidden): |
|
super(FCNet, self).__init__() |
|
self.inc_bias = False |
|
self.class_emb = nn.Linear(dim_hidden, num_classes, bias=self.inc_bias) |
|
|
|
self.feats = nn.Sequential(nn.Linear(num_inputs, dim_hidden), |
|
nn.ReLU(inplace=True), |
|
ResLayer(dim_hidden), |
|
ResLayer(dim_hidden), |
|
ResLayer(dim_hidden), |
|
ResLayer(dim_hidden)) |
|
|
|
def forward(self, x): |
|
loc_emb = self.feats(x) |
|
class_pred = self.class_emb(loc_emb) |
|
return class_pred |
|
|
|
"""A simple Multi Layer Perceptron""" |
|
class MLP(nn.Module): |
|
def __init__(self, input_dim, dim_hidden, num_layers, out_dims): |
|
super(MLP, self).__init__() |
|
|
|
layers = [] |
|
layers += [nn.Linear(input_dim, dim_hidden, bias=True), nn.ReLU()] |
|
layers += [nn.Linear(dim_hidden, dim_hidden, bias=True), nn.ReLU()] * num_layers |
|
layers += [nn.Linear(dim_hidden, out_dims, bias=True)] |
|
|
|
self.features = nn.Sequential(*layers) |
|
|
|
def forward(self, x): |
|
return self.features(x) |
|
|
|
def exists(val): |
|
return val is not None |
|
|
|
def cast_tuple(val, repeat = 1): |
|
return val if isinstance(val, tuple) else ((val,) * repeat) |
|
|
|
"""Sinusoidal Representation Network (SIREN)""" |
|
class SirenNet(nn.Module): |
|
def __init__(self, dim_in, dim_hidden, dim_out, num_layers, w0 = 1., w0_initial = 30., use_bias = True, final_activation = None, degreeinput = False, dropout = True): |
|
super().__init__() |
|
self.num_layers = num_layers |
|
self.dim_hidden = dim_hidden |
|
self.degreeinput = degreeinput |
|
|
|
self.layers = nn.ModuleList([]) |
|
for ind in range(num_layers): |
|
is_first = ind == 0 |
|
layer_w0 = w0_initial if is_first else w0 |
|
layer_dim_in = dim_in if is_first else dim_hidden |
|
|
|
self.layers.append(Siren( |
|
dim_in = layer_dim_in, |
|
dim_out = dim_hidden, |
|
w0 = layer_w0, |
|
use_bias = use_bias, |
|
is_first = is_first, |
|
dropout = dropout |
|
)) |
|
|
|
final_activation = nn.Identity() if not exists(final_activation) else final_activation |
|
self.last_layer = Siren(dim_in = dim_hidden, dim_out = dim_out, w0 = w0, use_bias = use_bias, activation = final_activation, dropout = False) |
|
|
|
def forward(self, x, mods = None): |
|
|
|
|
|
if self.degreeinput: |
|
x = torch.deg2rad(x) - torch.pi |
|
|
|
mods = cast_tuple(mods, self.num_layers) |
|
|
|
for layer, mod in zip(self.layers, mods): |
|
x = layer(x) |
|
|
|
if exists(mod): |
|
x *= rearrange(mod, 'd -> () d') |
|
|
|
return self.last_layer(x) |
|
|
|
class Sine(nn.Module): |
|
def __init__(self, w0 = 1.): |
|
super().__init__() |
|
self.w0 = w0 |
|
def forward(self, x): |
|
return torch.sin(self.w0 * x) |
|
|
|
class Siren(nn.Module): |
|
def __init__(self, dim_in, dim_out, w0 = 1., c = 6., is_first = False, use_bias = True, activation = None, dropout = False): |
|
super().__init__() |
|
self.dim_in = dim_in |
|
self.is_first = is_first |
|
self.dim_out = dim_out |
|
self.dropout = dropout |
|
|
|
weight = torch.zeros(dim_out, dim_in) |
|
bias = torch.zeros(dim_out) if use_bias else None |
|
self.init_(weight, bias, c = c, w0 = w0) |
|
|
|
self.weight = nn.Parameter(weight) |
|
self.bias = nn.Parameter(bias) if use_bias else None |
|
self.activation = Sine(w0) if activation is None else activation |
|
|
|
def init_(self, weight, bias, c, w0): |
|
dim = self.dim_in |
|
|
|
w_std = (1 / dim) if self.is_first else (math.sqrt(c / dim) / w0) |
|
weight.uniform_(-w_std, w_std) |
|
|
|
if exists(bias): |
|
bias.uniform_(-w_std, w_std) |
|
|
|
def forward(self, x): |
|
out = F.linear(x, self.weight, self.bias) |
|
if self.dropout: |
|
out = F.dropout(out, training=self.training) |
|
out = self.activation(out) |
|
return out |
|
|
|
|
|
class Modulator(nn.Module): |
|
def __init__(self, dim_in, dim_hidden, num_layers): |
|
super().__init__() |
|
self.layers = nn.ModuleList([]) |
|
|
|
for ind in range(num_layers): |
|
is_first = ind == 0 |
|
dim = dim_in if is_first else (dim_hidden + dim_in) |
|
|
|
self.layers.append(nn.Sequential( |
|
nn.Linear(dim, dim_hidden), |
|
nn.ReLU() |
|
)) |
|
|
|
def forward(self, z): |
|
x = z |
|
hiddens = [] |
|
|
|
for layer in self.layers: |
|
x = layer(x) |
|
hiddens.append(x) |
|
x = torch.cat((x, z)) |
|
|
|
return tuple(hiddens) |
|
|
|
class SirenWrapper(nn.Module): |
|
def __init__(self, net, image_width, image_height, latent_dim = None): |
|
super().__init__() |
|
assert isinstance(net, SirenNet), 'SirenWrapper must receive a Siren network' |
|
|
|
self.net = net |
|
self.image_width = image_width |
|
self.image_height = image_height |
|
|
|
self.modulator = None |
|
if exists(latent_dim): |
|
self.modulator = Modulator( |
|
dim_in = latent_dim, |
|
dim_hidden = net.dim_hidden, |
|
num_layers = net.num_layers |
|
) |
|
|
|
tensors = [torch.linspace(-1, 1, steps = image_height), torch.linspace(-1, 1, steps = image_width)] |
|
mgrid = torch.stack(torch.meshgrid(*tensors, indexing = 'ij'), dim=-1) |
|
mgrid = rearrange(mgrid, 'h w c -> (h w) c') |
|
self.register_buffer('grid', mgrid) |
|
|
|
def forward(self, img = None, *, latent = None): |
|
modulate = exists(self.modulator) |
|
assert not (modulate ^ exists(latent)), 'latent vector must be only supplied if `latent_dim` was passed in on instantiation' |
|
|
|
mods = self.modulator(latent) if modulate else None |
|
|
|
coords = self.grid.clone().detach().requires_grad_() |
|
out = self.net(coords, mods) |
|
out = rearrange(out, '(h w) c -> () c h w', h = self.image_height, w = self.image_width) |
|
|
|
if exists(img): |
|
return F.mse_loss(img, out) |
|
|
|
return out |
|
|
|
def get_positional_encoding(name, legendre_polys=10, harmonics_calculation='analytic', min_radius=1, max_radius=360, frequency_num=10): |
|
if name == "direct": |
|
return PE.Direct() |
|
elif name == "cartesian3d": |
|
return PE.Cartesian3D() |
|
elif name == "sphericalharmonics": |
|
if harmonics_calculation == 'discretized': |
|
return PE.DiscretizedSphericalHarmonics(legendre_polys=legendre_polys) |
|
else: |
|
return PE.SphericalHarmonics(legendre_polys=legendre_polys, |
|
harmonics_calculation=harmonics_calculation) |
|
elif name == "theory": |
|
return PE.Theory(min_radius=min_radius, |
|
max_radius=max_radius, |
|
frequency_num=frequency_num) |
|
elif name == "wrap": |
|
return PE.Wrap() |
|
elif name in ["grid", "spherec", "spherecplus", "spherem", "spheremplus"]: |
|
return PE.GridAndSphere(min_radius=min_radius, |
|
max_radius=max_radius, |
|
frequency_num=frequency_num, |
|
name=name) |
|
else: |
|
raise ValueError(f"{name} not a known positional encoding.") |
|
|
|
def get_neural_network(name, input_dim, num_classes=256, dim_hidden=256, num_layers=2): |
|
if name == "linear": |
|
return nn.Linear(input_dim, num_classes) |
|
elif name == "mlp": |
|
return MLP( |
|
input_dim=input_dim, |
|
dim_hidden=dim_hidden, |
|
num_layers=num_layers, |
|
out_dims=num_classes |
|
) |
|
elif name == "siren": |
|
return SirenNet( |
|
dim_in=input_dim, |
|
dim_hidden=dim_hidden, |
|
num_layers=num_layers, |
|
dim_out=num_classes |
|
) |
|
elif name == "fcnet": |
|
return FCNet( |
|
num_inputs=input_dim, |
|
num_classes=num_classes, |
|
dim_hidden=dim_hidden |
|
) |
|
else: |
|
raise ValueError(f"{name} not a known neural networks.") |
|
|
|
class LocationEncoder(nn.Module): |
|
def __init__(self, posenc, nnet): |
|
super().__init__() |
|
self.posenc = posenc |
|
self.nnet = nnet |
|
|
|
def forward(self, x): |
|
x = self.posenc(x) |
|
return self.nnet(x) |