diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index 320012ddc..d86209067 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -16,16 +16,15 @@ jobs: steps: - uses: actions/checkout@v3 - - name: Set up Python 3.10 + - name: Set up Python 3.8.12 uses: actions/setup-python@v3 with: - python-version: "3.10" + python-version: "3.8.12" - name: Install dependencies run: | python -m pip install --upgrade pip - pip install flake8 pytest - if [ -f environment/requirements.txt ]; then pip install -r environment/requirements.txt; fi + pip install -r environment/requirements.txt python setup.py develop - name: Test with pytest run: | - pytest tests/nerf/test_graph.py + pytest tests diff --git a/environment/requirements.txt b/environment/requirements.txt index ef64b7b0c..0c43402bc 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -4,6 +4,4 @@ pytest torch==1.11.0 torchtyping torchvision==0.12.0 -torchaudio==0.11.0 black -flake8 \ No newline at end of file diff --git a/mattport/nerf/modules/encoding.py b/mattport/nerf/modules/encoding.py index b0985ecfc..53ce14854 100644 --- a/mattport/nerf/modules/encoding.py +++ b/mattport/nerf/modules/encoding.py @@ -1,6 +1,7 @@ """ Encoding functions """ +import torch from torch import nn from torchtyping import TensorType @@ -8,9 +9,23 @@ from torchtyping import TensorType class Encoding(nn.Module): """Encode an input tensor. Intended to be subclassed""" - def __init__(self, out_dim: int) -> None: + def __init__(self, in_dim: int) -> None: + """ + Args: + in_dim (int): Input dimension of tensor + """ super().__init__() - self.out_dim = out_dim + if in_dim <= 0: + raise ValueError("Input dimension should be greater than zero") + self.in_dim = in_dim + + def get_out_dim(self) -> int: + """Calculates output dimension of encoding. + + Returns: + int: output dimension + """ + raise NotImplementedError def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: """Encodes an input tensor. @@ -28,19 +43,98 @@ class Encoding(nn.Module): return self.encode(in_tensor) +class ScalingAndOffset(Encoding): + """Simple scaling and offet to input""" + + def __init__(self, in_dim: int, scaling: float = 1.0, offset: float = 0.0) -> None: + """Each input is scaled then offset + + Args: + in_dim (int): Input dimension of tensor + scaling (float, optional): Scaling applied to tensor. Defaults to 1.0. + offset (float, optional): Offset applied to tensor. Defaults to 0.0. + """ + super().__init__(in_dim) + + self.scaling = scaling + self.offset = offset + + def get_out_dim(self) -> int: + return self.in_dim + + def encode(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""" + + def __init__(self, in_dim: int, num_frequencies: int, min_freq_exp: float, max_freq_exp: float) -> None: + """Each axis is encoded with frequencies ranging from 2^min_freq_exp to 2^max_freq_exp. + + Args: + in_dim (int): Input dimension of tensor + num_frequencies (int): Number of encoded frequencies per axis + min_freq_exp (float): Minimum frequency exponent + max_freq_exp (float): Maximum frequency exponent + """ + super().__init__(in_dim) + + self.num_frequencies = num_frequencies + self.min_freq = min_freq_exp + self.max_freq = max_freq_exp + + def get_out_dim(self) -> int: + return self.in_dim * self.num_frequencies * 2 + + def encode(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. + + Returns: + TensorType[..., "output_dim"]: Output values will be between -1 and 1 + """ + freqs = 2 ** torch.linspace(self.min_freq, self.max_freq, self.num_frequencies) + scaled_inputs = 2 * torch.pi * 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) + return encoded_inputs + + class RFFEncoding(Encoding): """Random Fourier Feature encoding""" - def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: - """Encodes an input tensor. + def __init__(self, in_dim: int, num_frequencies: int, scale: float) -> None: + """ Args: - x (TensorType[..., "input_dim"]): Input tensor to be encoded + in_dim (int): Input dimension of tensor + num_frequencies (int): Number of encoding frequencies + scale (float): Std of Gaussian to sample frequencies. Must be greater than zero + """ + super().__init__(in_dim) + + self.num_frequencies = num_frequencies + if not scale > 0: + raise ValueError("RFF encoding scale should be greater than zero") + self.scale = scale + + def get_out_dim(self) -> int: + return self.num_frequencies * 2 + + def encode(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. Returns: - TensorType[..., "output_dim"]: A encoded input tensor + TensorType[..., "output_dim"]: Output values will be between -1 and 1 """ - raise NotImplementedError + B = torch.normal(mean=0, std=self.scale, size=(self.in_dim, self.num_frequencies)) + scaled_inputs = in_tensor @ B # [..., "num_frequencies"] + encoded_inputs = torch.cat([torch.sin(scaled_inputs), torch.cos(scaled_inputs)], axis=-1) + return encoded_inputs class HashEncoding(Encoding): @@ -56,9 +150,3 @@ class HashEncoding(Encoding): TensorType[..., "output_dim"]: A encoded input tensor """ raise NotImplementedError - - -encodings_dict = { - "RFF": RFFEncoding, - "hash": HashEncoding, -} diff --git a/tests/nerf/modules/test_encodings.py b/tests/nerf/modules/test_encodings.py new file mode 100644 index 000000000..47bd74668 --- /dev/null +++ b/tests/nerf/modules/test_encodings.py @@ -0,0 +1,65 @@ +""" +Encoding Tests +""" +import pytest +import torch +from mattport.nerf.modules import encoding + + +def test_scaling_and_offset(): + """Test scaling and offset encoder""" + in_dim = 4 + in_tensor = torch.ones((2, 3, in_dim)) + + scaling = 2.0 + offset = 4.5 + encoder = encoding.ScalingAndOffset(in_dim=in_dim, scaling=scaling, offset=offset) + + assert encoder.get_out_dim() == in_dim + encoded = encoder.encode(in_tensor) + assert encoded.shape[-1] == in_dim + assert in_tensor * 6.5 == pytest.approx(encoded) + + with pytest.raises(ValueError): + encoding.ScalingAndOffset(in_dim=-1) + + +def test_nerf_encoder(): + """Test NeRF encoder""" + in_dim = 4 + out_dim = 24 + in_tensor = torch.ones((2, 3, in_dim)) + + num_frequencies = 3 + min_freq_exp = 0 + max_freq_exp = 3 + encoder = encoding.NeRFEncoding( + in_dim=in_dim, num_frequencies=num_frequencies, min_freq_exp=min_freq_exp, max_freq_exp=max_freq_exp + ) + assert encoder.get_out_dim() == out_dim + + in_tensor = torch.ones((2, 3, in_dim)) + encoded = encoder.encode(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) + assert encoded.shape[-1] == out_dim + assert torch.min(encoded) == 0 + + +def test_rff_encoder(): + """Test RFF encoder""" + in_dim = 3 + out_dim = 24 + in_tensor = torch.ones((2, 3, in_dim)) + + num_frequencies = 12 + scale = 5 + encoder = encoding.RFFEncoding(in_dim=in_dim, num_frequencies=num_frequencies, scale=scale) + assert encoder.get_out_dim() == out_dim + + in_tensor = torch.ones((2, 3, in_dim)) + encoded = encoder.encode(in_tensor) + assert encoded.shape[-1] == out_dim diff --git a/tests/nerf/test_encodings.py b/tests/nerf/test_encodings.py deleted file mode 100644 index 99f26096c..000000000 --- a/tests/nerf/test_encodings.py +++ /dev/null @@ -1,9 +0,0 @@ -import pytest - - -def test_positional_encoding(): - assert True - - -def test_lower_case(): - assert "MATTPORT".lower() == "mattport"