Triplanes (#1456)

* Triplanes working

* Doing the TODOs, better variable names

---------

Co-authored-by: Matthew Tancik <tancik@berkeley.edu>

nerfstudio/field_components/encodings.py

@@ -469,11 +469,94 @@ class TensorVMEncoding(Encoding):
         )
         line_coef = F.interpolate(self.line_coef.data, size=(resolution, 1), mode="bilinear", align_corners=True)
 
-        # TODO(ethan): are these torch.nn.Parameters needed?
         self.plane_coef, self.line_coef = torch.nn.Parameter(plane_coef), torch.nn.Parameter(line_coef)
         self.resolution = resolution
 
 
+class TriplaneEncoding(Encoding):
+    """Learned triplane encoding
+
+    The encoding at [i,j,k] is an n dimensional vector corresponding to the element-wise product of the
+    three n dimensional vectors at plane_coeff[i,j], plane_coeff[i,k], and plane_coeff[j,k].
+
+    This allows for marginally more expressivity than the TensorVMEncoding, and each component is self standing
+    and symmetrical, unlike with VM decomposition where we needed one component with a vector along all the x, y, z
+    directions for symmetry.
+
+    This can be thought of as 3 planes of features perpendicular to the x, y, and z axes, respectively and intersecting
+    at the origin, and the encoding being the element-wise product of the element at the projection of [i, j, k] on
+    these planes.
+
+    The use for this is in representing a tensor decomp of a 4D embedding tensor: (x, y, z, feature_size)
+
+    This will return a tensor of shape (bs:..., num_components)
+
+    Args:
+        resolution: Resolution of grid.
+        num_components: The number of scalar triplanes to use (ie: output feature size)
+        init_scale: The scale of the initial values of the planes
+        product: Whether to use the element-wise product of the planes or the sum
+    """
+
+    plane_coef: TensorType[3, "num_components", "resolution", "resolution"]
+
+    def __init__(
+        self,
+        resolution: int = 32,
+        num_components: int = 64,
+        init_scale: float = 0.1,
+        reduce: Literal["sum", "product"] = "sum",
+    ) -> None:
+        super().__init__(in_dim=3)
+
+        self.resolution = resolution
+        self.num_components = num_components
+        self.init_scale = init_scale
+        self.reduce = reduce
+
+        self.plane_coef = nn.Parameter(
+            self.init_scale * torch.randn((3, self.num_components, self.resolution, self.resolution))
+        )
+
+    def get_out_dim(self) -> int:
+        return self.num_components
+
+    def forward(self, in_tensor: TensorType["bs":..., 3]) -> TensorType["bs":..., "num_components", "featuresize"]:
+        """Sample features from this encoder. Expects in_tensor to be in range [0, resolution]"""
+
+        original_shape = in_tensor.shape
+        in_tensor = in_tensor.reshape(-1, 3)
+
+        plane_coord = torch.stack([in_tensor[..., [0, 1]], in_tensor[..., [0, 2]], in_tensor[..., [1, 2]]], dim=0)
+
+        # Stop gradients from going to sampler
+        plane_coord = plane_coord.detach().view(3, -1, 1, 2)
+        plane_features = F.grid_sample(
+            self.plane_coef, plane_coord, align_corners=True
+        )  # [3, num_components, flattened_bs, 1]
+
+        if self.reduce == "product":
+            plane_features = plane_features.prod(0).squeeze(-1).T  # [flattened_bs, num_components]
+        else:
+            plane_features = plane_features.sum(0).squeeze(-1).T
+
+        return plane_features.reshape(*original_shape[:-1], self.num_components)
+
+    @torch.no_grad()
+    def upsample_grid(self, resolution: int) -> None:
+        """Upsamples underlying feature grid
+
+        Args:
+            resolution: Target resolution.
+        """
+        plane_coef = F.interpolate(
+            self.plane_coef.data, size=(resolution, resolution), mode="bilinear", align_corners=True
+        )
+
+        self.plane_coef = torch.nn.Parameter(plane_coef)
+        self.resolution = resolution
+
+
 class SHEncoding(Encoding):
     """Spherical harmonic encoding
 

nerfstudio/models/tensorf.py

@@ -27,6 +27,7 @@ from torch.nn import Parameter
 from torchmetrics import PeakSignalNoiseRatio
 from torchmetrics.functional import structural_similarity_index_measure
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from typing_extensions import Literal
 
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.configs.config_utils import to_immutable_dict
@@ -35,7 +36,12 @@ from nerfstudio.engine.callbacks import (
     TrainingCallbackAttributes,
     TrainingCallbackLocation,
 )
-from nerfstudio.field_components.encodings import NeRFEncoding, TensorVMEncoding
+from nerfstudio.field_components.encodings import (
+    NeRFEncoding,
+    TensorCPEncoding,
+    TensorVMEncoding,
+    TriplaneEncoding,
+)
 from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.fields.tensorf_field import TensoRFField
 from nerfstudio.model_components.losses import MSELoss
@@ -72,6 +78,7 @@ class TensoRFModelConfig(ModelConfig):
     """Number of components in color encoding"""
     appearance_dim: int = 27
     """Number of channels for color encoding"""
+    tensorf_encoding: Literal["triplane", "vm", "cp"] = "vm"
 
 
 class TensoRFModel(Model):
@@ -166,14 +173,35 @@ class TensoRFModel(Model):
         super().populate_modules()
 
         # setting up fields
-        density_encoding = TensorVMEncoding(
-            resolution=self.init_resolution,
-            num_components=self.num_den_components,
-        )
-        color_encoding = TensorVMEncoding(
-            resolution=self.init_resolution,
-            num_components=self.num_color_components,
-        )
+        if self.config.tensorf_encoding == "vm":
+            density_encoding = TensorVMEncoding(
+                resolution=self.init_resolution,
+                num_components=self.num_den_components,
+            )
+            color_encoding = TensorVMEncoding(
+                resolution=self.init_resolution,
+                num_components=self.num_color_components,
+            )
+        elif self.config.tensorf_encoding == "cp":
+            density_encoding = TensorCPEncoding(
+                resolution=self.init_resolution,
+                num_components=self.num_den_components,
+            )
+            color_encoding = TensorCPEncoding(
+                resolution=self.init_resolution,
+                num_components=self.num_color_components,
+            )
+        elif self.config.tensorf_encoding == "triplane":
+            density_encoding = TriplaneEncoding(
+                resolution=self.init_resolution,
+                num_components=self.num_den_components,
+            )
+            color_encoding = TriplaneEncoding(
+                resolution=self.init_resolution,
+                num_components=self.num_color_components,
+            )
+        else:
+            raise ValueError(f"Encoding {self.config.tensorf_encoding} not supported")
 
         feature_encoding = NeRFEncoding(in_dim=self.appearance_dim, num_frequencies=2, min_freq_exp=0, max_freq_exp=2)
         direction_encoding = NeRFEncoding(in_dim=3, num_frequencies=2, min_freq_exp=0, max_freq_exp=2)