diff --git a/nerfstudio/field_components/encodings.py b/nerfstudio/field_components/encodings.py index 9cdd85e61..f77c17f4d 100644 --- a/nerfstudio/field_components/encodings.py +++ b/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 diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index 48310fde2..3b27354d1 100644 --- a/nerfstudio/models/tensorf.py +++ b/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)