File size: 1,619 Bytes
1eb87a5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import math
import jax
import jax.numpy as jnp
import flax.linen as nn
from jaxtyping import Array, ArrayLike, PyTreeDef
import numpy as np
from utils import interpolate_grid
class Hypernetwork(nn.Module):
encoder: nn.Module
refine: nn.Module
output_params_shape: list[tuple] # e.g. [(16,), (32, 32), ...]
tree_def: PyTreeDef # used to reconstruct the parameter sets
def setup(self):
# one layer 1x1 conv to calculate field params, as in SIREN paper
output_size = sum(math.prod(s) for s in self.output_params_shape)
self.out_conv = nn.Conv(output_size, kernel_size=(1, 1), use_bias=True)
def get_encoding(self, source: ArrayLike, training=False) -> Array:
"""Convenience method for whole-image evaluation"""
return self.refine(self.encoder(source, training), training)
def get_params_at_coords(self, encoding: ArrayLike, coords: ArrayLike) -> Array:
encoding = interpolate_grid(coords, encoding)
phi_params = self.out_conv(encoding)
# reshape to output params shape
phi_params = jnp.split(
phi_params, np.cumsum([math.prod(s) for s in self.output_params_shape[:-1]]), axis=-1)
phi_params = [jnp.reshape(p, p.shape[:-1] + s) for p, s in
zip(phi_params, self.output_params_shape)]
return jax.tree_util.tree_unflatten(self.tree_def, phi_params)
def __call__(self, source: ArrayLike, target_coords: ArrayLike, training=False) -> Array:
encoding = self.get_encoding(source, training)
return self.get_params_at_coords(encoding, target_coords)
|