Custom components¶
The library is built from small interfaces. This page shows the contract of each one with a minimal working example.
| Component | Built or called as | Must provide |
|---|---|---|
| Basis | basis(in_features, **kwargs) |
num_components, and one (N, K_c) feature tensor per axis |
| Combiner | combiner(basis.num_components, **kwargs) |
out_features, and (N, out_features) from the features |
| Head | head(combiner.out_features, out_features, **kwargs) |
Any module |
| Loss | loss_fn(batch, model) |
A scalar loss and a dictionary of metrics |
| Metric | fn(reconstruction, original) |
A float |
| Dataset | dataset[i] |
Items with id, input, target, _original; preprocess, postprocess, save |
| Model | model(coordinates) |
(..., in_features) to (..., out_features) |
A basis¶
A Gaussian bump basis in a dozen lines:
import torch
from torch import nn
import inro
class GaussianBasis(nn.Module):
"""K Gaussian bumps per axis, centered on a uniform grid over [-1, 1]."""
def __init__(self, in_features, num_components, width=0.05):
super().__init__()
self.in_features = in_features
self.num_components = [num_components] * in_features
self.width = width
self.register_buffer("centers", torch.linspace(-1.0, 1.0, num_components))
def forward(self, x): # (N, C) -> C tensors of shape (N, K)
return tuple(
torch.exp(-(((x[:, c, None] - self.centers) / self.width) ** 2))
for c in range(self.in_features)
)
model = inro.FUTON(
2, 3,
basis=(GaussianBasis, {"num_components": 64}),
combiner=("cp", {"rank": 32}),
head="linear",
)
A class and a (class, kwargs) pair work anywhere a registry key does, as does a ready module instance. To add a key instead, register the class in inro.models.futon.BASES.
A combiner or head¶
A combiner takes the tuple of per-axis features: tensors, SharedRows on a grid, or Taps for a local basis. A product with a matrix gives a tensor, or SharedRows again on a grid, whose dense() is the matrix. A head is any module, so torch.nn.Linear, inro.MLP, or a class of your own:
class GatedHead(nn.Module):
def __init__(self, in_features, out_features):
super().__init__()
self.linear = nn.Linear(in_features, out_features)
self.gate = nn.Linear(in_features, out_features)
def forward(self, g):
return self.linear(g) * torch.sigmoid(self.gate(g))
model = inro.FUTON(
2, 3,
basis=("cosine", {"num_components": 128}),
combiner=("cp", {"rank": 64}),
head=GatedHead,
)
A loss¶
inro.train calls loss_fn(batch, model) and expects a scalar loss and a dictionary of metrics for the log, as floats or as scalar tensors. Tensors are read only when a step is logged, so the training steps never wait for the GPU in between. batch is the dataset item with its tensors on the training device; keys with a leading underscore stay on the CPU.
A loss that fits only the pixels a mask marks as known:
import torch.nn.functional as F
def masked_mse(batch, model):
output = model(batch["input"])
mask = batch["mask"]
loss = F.mse_loss(output[mask], batch["target"][mask])
return loss, {"loss": loss.detach()}
train_set = inro.MaskedImageCoordinateDataset(
"image.png", mask="mask.png", subsample=0.1
)
eval_set = inro.ImageCoordinateDataset("image.png")
result = inro.train(model, train_set, eval_set, loss_fn=masked_mse)
The indexing works before and after subsampling because the dataset flattens input, target and mask with the same indices. inro.rate_distortion_loss and inro.sdf_loss are two more examples of the signature; functools.partial(inro.rate_distortion_loss, alpha=0.01) sets their options.
A metric¶
Metrics are functions fn(reconstruction, original) -> float on whatever the dataset's postprocess returns and its _original holds: uint8 (channels, H, W) images or {0, 1} volumes. Rate metrics take (encoded_bytes, original).
def mae(pred, target):
return (pred.float() - target.float()).abs().mean().item()
result = inro.train(model, train_set, eval_set, metrics={"psnr": inro.psnr, "mae": mae})
For inro.nerf.evaluate, metrics receive (3, H, W) float images in \([0, 1]\).
A dataset¶
A dataset is a torch.utils.data.Dataset whose items carry id, input, target and _original, with preprocess, postprocess and save for the evaluation loop and, optionally, to(device) for GPU-side subsampling. The simplest route is to subclass ImageCoordinateDataset and override what differs; OccupancyCoordinateDataset and MaskedImageCoordinateDataset do exactly that. For a signal with its own geometry, build the coordinates with inro.create_coordinates(size) and follow the conventions.
A model¶
Any module that maps (..., in_features) coordinates in \([-1, 1]\) to (..., out_features) is an INR for the training loop, and, without an output activation, a density network for inro.nerf.RadianceField. inro.MLP accepts a custom layer_class when the novelty is in the layer:
class SoftplusLayer(nn.Module):
def __init__(self, in_features, out_features, beta=1.0):
super().__init__()
self.linear = nn.Linear(in_features, out_features)
self.beta = beta
def forward(self, x):
return F.softplus(self.linear(x), beta=self.beta)
model = inro.MLP(
2, 3, hidden_features=256, hidden_layers=3,
layer_class=SoftplusLayer, beta=2.0, output_activation=torch.tanh,
)
Extra keyword arguments go to the hidden layers; the output layer stays linear. SIREN, FINER, Gauss and RealWIRE are built this way and add their initialization after super().__init__().
Registries¶
inro.build_module(spec, registry, *args, **kwargs) is the resolver behind every spec: an instance is returned as is, a key is looked up in registry, a class is called, and a (key_or_class, params) pair adds keyword arguments. A two-element list works too, so specs can come from YAML or JSON. Use it for your own components: