Merge pull request #10 from ethanweber/tancik/implement_encodings

Tancik/implement encodings

.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

environment/requirements.txt

@@ -4,6 +4,4 @@ pytest
 torch==1.11.0
 torchtyping
 torchvision==0.12.0
-torchaudio==0.11.0
 black
-flake8

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,
-}

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

tests/nerf/test_encodings.py

@@ -1,9 +0,0 @@
-import pytest
-
-
-def test_positional_encoding():
-    assert True
-
-
-def test_lower_case():
-    assert "MATTPORT".lower() == "mattport"