From cc4d3431b831316fa01bddbe9a6c62ac26d2e55a Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Wed, 8 Jun 2022 16:09:46 -0700 Subject: [PATCH] Add integrated positional encodings --- pyrad/fields/modules/base.py | 10 ++--- pyrad/fields/modules/encoding.py | 61 +++++++++++++------------- pyrad/utils/math.py | 16 ++++++- tests/fields/modules/test_encodings.py | 18 ++++---- 4 files changed, 59 insertions(+), 46 deletions(-) diff --git a/pyrad/fields/modules/base.py b/pyrad/fields/modules/base.py index 6009b9e64..5adae1e8e 100644 --- a/pyrad/fields/modules/base.py +++ b/pyrad/fields/modules/base.py @@ -24,11 +24,11 @@ from torchtyping import TensorType class FieldModule(nn.Module): """Field modules that can be combined to store and compute the fields.""" - def __init__(self) -> None: + def __init__(self, in_dim=None, out_dim=None) -> None: """Default initialization of module""" super().__init__() - self.in_dim = 0 - self.out_dim = 0 + self.in_dim = in_dim + self.out_dim = out_dim def build_nn_modules(self) -> None: """Function instantiates any torch.nn members within the module. @@ -50,10 +50,8 @@ class FieldModule(nn.Module): Returns: int: output dimension """ - if not hasattr(self, "out_dim"): + if self.out_dim is None: raise ValueError("Output dimension has not been set") - if self.out_dim <= 0: - raise ValueError("Output dimension should be greater than zero") return self.out_dim @abstractmethod diff --git a/pyrad/fields/modules/encoding.py b/pyrad/fields/modules/encoding.py index c7a068fa6..577c308ba 100644 --- a/pyrad/fields/modules/encoding.py +++ b/pyrad/fields/modules/encoding.py @@ -17,6 +17,7 @@ Encoding functions """ from abc import abstractmethod +from typing import Optional import numpy as np import torch import torch.nn.functional as F @@ -24,7 +25,7 @@ from torch import nn from torchtyping import TensorType from pyrad.fields.modules.base import FieldModule -from pyrad.utils.math import components_from_spherical_harmonics +from pyrad.utils.math import components_from_spherical_harmonics, expected_sin class Encoding(FieldModule): @@ -35,26 +36,13 @@ class Encoding(FieldModule): Args: in_dim (int): Input dimension of tensor """ - super().__init__() if in_dim <= 0: raise ValueError("Input dimension should be greater than zero") - self.in_dim = in_dim + super().__init__(in_dim=in_dim) @abstractmethod - def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: - """Encodes an input tensor. - - Args: - x (TensorType[..., "input_dim"]): Input tensor to be encoded - - Returns: - TensorType[..., "output_dim"]: A encoded input tensor - """ - raise NotImplementedError - def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: - """Call forward""" - return self.encode(in_tensor) + raise NotImplementedError class Identity(Encoding): @@ -63,7 +51,7 @@ class Identity(Encoding): def get_out_dim(self) -> int: return self.in_dim - def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: return in_tensor @@ -86,12 +74,12 @@ class ScalingAndOffset(Encoding): def get_out_dim(self) -> int: return self.in_dim - def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: return self.scaling * in_tensor + self.offset class NeRFEncoding(Encoding): - """Multi-scale sinousoidal encoding proposed in the original NeRF paper""" + """Multi-scale sinousoidal encodings. Support ``integrated positional encodings`` if covariances are provided.""" def __init__( self, in_dim: int, num_frequencies: int, min_freq_exp: float, max_freq_exp: float, include_input: bool = False @@ -118,18 +106,32 @@ class NeRFEncoding(Encoding): out_dim += self.in_dim return out_dim - def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: - """ + def forward( + self, in_tensor: TensorType[..., "input_dim"], covs: Optional[TensorType[..., "input_dim", "input_dim"]] = None + ) -> TensorType[..., "output_dim"]: + """Calculates NeRF encoding. If covariances are provided the encodings will be integrated as proposed + in mip-NeRF. + Args: in_tensor (TensorType[..., "input_dim"]): For best performance, the input tensor should be between 0 and 1. - + covs (TensorType[..., "input_dim", "input_dim"], optional): Covariances of input points. Defaults to None. Returns: TensorType[..., "output_dim"]: Output values will be between -1 and 1 """ + in_tensor = 2 * torch.pi * in_tensor # scale to [0, 2pi] freqs = 2 ** torch.linspace(self.min_freq, self.max_freq, self.num_frequencies).to(in_tensor.device) - scaled_inputs = 2 * torch.pi * in_tensor[..., None] * freqs # [..., "input_dim", "num_scales"] + scaled_inputs = in_tensor[..., None] * freqs # [..., "input_dim", "num_scales"] scaled_inputs = scaled_inputs.view(*scaled_inputs.shape[:-2], -1) # [..., "input_dim" * "num_scales"] - encoded_inputs = torch.cat([torch.sin(scaled_inputs), torch.cos(scaled_inputs)], axis=-1) + + if covs is None: + encoded_inputs = torch.sin(torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1)) + else: + input_var = torch.diagonal(covs, dim1=-2, dim2=-1)[..., None, :] * freqs[:, None] ** 2 + input_var = input_var.reshape((*input_var.shape[:-2], -1)) + encoded_inputs = expected_sin( + torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1), torch.cat(2 * [input_var], dim=-1) + ) + if self.include_input: encoded_inputs = torch.cat([encoded_inputs, in_tensor], axis=-1) return encoded_inputs @@ -158,7 +160,7 @@ class RFFEncoding(Encoding): def get_out_dim(self) -> int: return self.num_frequencies * 2 - def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: """ Args: in_tensor (TensorType[..., "input_dim"]): For best performance, the input tensor should be between 0 and 1. @@ -232,7 +234,7 @@ class HashEncoding(Encoding): x += self.hash_offset.to(x.device) return x - def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: assert in_tensor.shape[-1] == 3 in_tensor = in_tensor[..., None, :] # [..., 1, 3] @@ -296,8 +298,7 @@ class TensorCPEncoding(Encoding): def get_out_dim(self) -> int: return self.num_components - def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: - in_tensor = in_tensor / 4.0 + 0.5 + def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: line_coord = torch.stack([in_tensor[..., 2], in_tensor[..., 1], in_tensor[..., 0]]) # [3, ...] line_coord = torch.stack([torch.zeros_like(line_coord), line_coord], dim=-1) # [3, ...., 2] @@ -348,7 +349,7 @@ class TensorVMEncoding(Encoding): def get_out_dim(self) -> int: return self.num_components * 3 - def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: in_tensor = in_tensor / 4.0 + 0.5 plane_coord = torch.stack([in_tensor[..., [0, 1]], in_tensor[..., [0, 2]], in_tensor[..., [1, 2]]]) # [3,...,2] line_coord = torch.stack([in_tensor[..., 2], in_tensor[..., 1], in_tensor[..., 0]]) # [3, ...] @@ -403,5 +404,5 @@ class SHEncoding(Encoding): return self.levels**2 @torch.no_grad() - def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: return components_from_spherical_harmonics(levels=self.levels, directions=in_tensor) diff --git a/pyrad/utils/math.py b/pyrad/utils/math.py index 97ffccc6a..66f9a8b3e 100644 --- a/pyrad/utils/math.py +++ b/pyrad/utils/math.py @@ -116,7 +116,7 @@ def compute_3d_gaussian( """ dir_outer_product = directions[..., :, None] * directions[..., None, :] - eye = torch.eye(directions.shape[-1]) + eye = torch.eye(directions.shape[-1], device=directions.device) dir_mag_sq = torch.clamp(torch.sum(directions**2, axis=-1, keepdim=True), min=1e-10) null_outer_product = eye - directions[..., :, None] * (directions / dir_mag_sq)[..., None, :] dir_cov_diag = dir_variance[..., None] * dir_outer_product[..., :, :] @@ -177,3 +177,17 @@ def conical_frustum_to_gaussian( dir_variance = (hw**2) / 3 - (4 / 15) * ((hw**4 * (12 * mu**2 - hw**2)) / (3 * mu**2 + hw**2) ** 2) radius_variance = radius**2 * ((mu**2) / 4 + (5 / 12) * hw**2 - 4 / 15 * (hw**4) / (3 * mu**2 + hw**2)) return compute_3d_gaussian(directions, means, dir_variance, radius_variance) + + +def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor: + """Computes the expected value of sin(y) where y ~ N(x_means, x_vars) + + Args: + x_means (torch.Tensor): Mean values. + x_vars (torch.Tensor): Variance of values. + + Returns: + torch.Tensor: The expected value of sin. + """ + + return torch.exp(-0.5 * x_vars) * torch.sin(x_means) diff --git a/tests/fields/modules/test_encodings.py b/tests/fields/modules/test_encodings.py index ff9c2965c..09c5caaa8 100644 --- a/tests/fields/modules/test_encodings.py +++ b/tests/fields/modules/test_encodings.py @@ -16,7 +16,7 @@ def test_scaling_and_offset(): encoder = encoding.ScalingAndOffset(in_dim=in_dim, scaling=scaling, offset=offset) assert encoder.get_out_dim() == in_dim - encoded = encoder.encode(in_tensor) + encoded = encoder(in_tensor) assert encoded.shape[-1] == in_dim assert in_tensor * 6.5 == pytest.approx(encoded) @@ -39,12 +39,12 @@ def test_nerf_encoder(): assert encoder.get_out_dim() == out_dim in_tensor = torch.ones((2, 3, in_dim)) - encoded = encoder.encode(in_tensor) + encoded = encoder(in_tensor) assert encoded.shape[-1] == out_dim assert torch.max(encoded) == 1 in_tensor = torch.zeros((2, 3, in_dim)) - encoded = encoder.encode(in_tensor) + encoded = encoder(in_tensor) assert encoded.shape[-1] == out_dim assert torch.min(encoded) == 0 @@ -61,7 +61,7 @@ def test_rff_encoder(): assert encoder.get_out_dim() == out_dim in_tensor = torch.ones((2, 3, in_dim)) - encoded = encoder.encode(in_tensor) + encoded = encoder(in_tensor) assert encoded.shape[-1] == out_dim @@ -78,11 +78,11 @@ def test_tensor_vm_encoder(): assert encoder.get_out_dim() == out_dim in_tensor = torch.ones((3, in_dim)) - encoded = encoder.encode(in_tensor) + encoded = encoder(in_tensor) assert encoded.shape == (3, out_dim) in_tensor = torch.ones((6, 3, in_dim)) - encoded = encoder.encode(in_tensor) + encoded = encoder(in_tensor) assert encoded.shape == (6, 3, out_dim) encoder.upsample_grid(resolution=64) @@ -101,11 +101,11 @@ def test_tensor_cp_encoder(): assert encoder.get_out_dim() == out_dim in_tensor = torch.ones((3, in_dim)) - encoded = encoder.encode(in_tensor) + encoded = encoder(in_tensor) assert encoded.shape == (3, out_dim) in_tensor = torch.ones((6, 3, in_dim)) - encoded = encoder.encode(in_tensor) + encoded = encoder(in_tensor) assert encoded.shape == (6, 3, out_dim) encoder.upsample_grid(resolution=64) @@ -125,7 +125,7 @@ def test_tensor_sh_encoder(): in_tensor = torch.zeros((10, 3)) in_tensor[..., 1] = 1 - encoded = encoder.encode(in_tensor) + encoded = encoder(in_tensor) assert encoded.shape == (10, out_dim)