From 7106f42fcc67456009502eccbeac96a2c5d2b623 Mon Sep 17 00:00:00 2001 From: Terrance Wang Date: Wed, 2 Nov 2022 14:32:03 -0700 Subject: [PATCH] tensorf (#586) * tensorf wip * model changes * tuning model * changing imports * merging * changing rff encoding to nerfencoding, basis matrix should be properly updated, changed density activation, lr scheduler * debugging tensorf * debugging tensorf * tensorf debugging * tensorf debugging * tensorf debugging * tensorf debug * debugging tensorf, fixed bug with collider * tuning model, ray sampling * adding thresholding for ray computations * removing unused imports * making more params configurable in config, adding spherical harmonics * removing scene contraction, adding arg docstrings * fixing style * Fixing bug Co-authored-by: Matthew Tancik --- nerfstudio/configs/method_configs.py | 29 ++- nerfstudio/field_components/encodings.py | 3 +- nerfstudio/fields/tensorf_field.py | 135 +++++++++++ nerfstudio/models/tensorf.py | 280 +++++++++++++++++++++++ 4 files changed, 444 insertions(+), 3 deletions(-) create mode 100644 nerfstudio/fields/tensorf_field.py create mode 100644 nerfstudio/models/tensorf.py diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index f4a95f027..f99110569 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -23,7 +23,12 @@ from typing import Dict import tyro from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig -from nerfstudio.configs.base_config import Config, TrainerConfig, ViewerConfig +from nerfstudio.configs.base_config import ( + Config, + SchedulerConfig, + TrainerConfig, + ViewerConfig, +) from nerfstudio.data.datamanagers import VanillaDataManagerConfig from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig @@ -34,6 +39,7 @@ from nerfstudio.models.instant_ngp import InstantNGPModelConfig from nerfstudio.models.mipnerf import MipNerfModel from nerfstudio.models.nerfacto import NerfactoModelConfig from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig +from nerfstudio.models.tensorf import TensoRFModelConfig from nerfstudio.models.vanilla_nerf import NeRFModel from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig @@ -45,6 +51,7 @@ descriptions = { "mipnerf": "High quality model for bounded scenes. (slow)", "semantic-nerfw": "Predicts semantic segmentations and filters out transient objects.", "vanilla-nerf": "Original NeRF model. (slow)", + "tensorf": "tensorf", } method_configs["nerfacto"] = Config( @@ -155,6 +162,26 @@ method_configs["vanilla-nerf"] = Config( }, ) +method_configs["tensorf"] = Config( + method_name="tensorf", + trainer=TrainerConfig(mixed_precision=False), + pipeline=VanillaPipelineConfig( + datamanager=VanillaDataManagerConfig( + dataparser=BlenderDataParserConfig(), + ), + model=TensoRFModelConfig(), + ), + optimizers={ + "fields": { + "optimizer": AdamOptimizerConfig(lr=0.001), + "scheduler": SchedulerConfig(lr_final=0.0001, max_steps=30000), + }, + "encodings": { + "optimizer": AdamOptimizerConfig(lr=0.02), + "scheduler": SchedulerConfig(lr_final=0.002, max_steps=30000), + }, + }, +) AnnotatedBaseConfigUnion = tyro.conf.SuppressFixed[ # Don't show unparseable (fixed) arguments in helptext. tyro.extras.subcommand_type_from_defaults(defaults=method_configs, descriptions=descriptions) diff --git a/nerfstudio/field_components/encodings.py b/nerfstudio/field_components/encodings.py index 7e149e531..65bff8244 100644 --- a/nerfstudio/field_components/encodings.py +++ b/nerfstudio/field_components/encodings.py @@ -416,7 +416,6 @@ class TensorVMEncoding(Encoding): self.resolution = resolution self.num_components = num_components - # TODO(terrance): Learning rates should be different for these self.plane_coef = nn.Parameter(init_scale * torch.randn((3, num_components, resolution, resolution))) self.line_coef = nn.Parameter(init_scale * torch.randn((3, num_components, resolution, 1))) @@ -424,7 +423,6 @@ class TensorVMEncoding(Encoding): return self.num_components * 3 def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "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, ...] line_coord = torch.stack([torch.zeros_like(line_coord), line_coord], dim=-1) # [3, ...., 2] @@ -452,6 +450,7 @@ class TensorVMEncoding(Encoding): self.plane_coef.data, size=(resolution, resolution), mode="bilinear", align_corners=True ) 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 diff --git a/nerfstudio/fields/tensorf_field.py b/nerfstudio/fields/tensorf_field.py new file mode 100644 index 000000000..d5cc9ba3b --- /dev/null +++ b/nerfstudio/fields/tensorf_field.py @@ -0,0 +1,135 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TensoRF Field""" + + +from typing import Optional + +import torch +from torch import nn +from torch.nn.parameter import Parameter +from torchtyping import TensorType + +from nerfstudio.cameras.rays import RaySamples +from nerfstudio.data.scene_box import SceneBox +from nerfstudio.field_components.encodings import Encoding, Identity, SHEncoding +from nerfstudio.field_components.field_heads import FieldHeadNames, RGBFieldHead +from nerfstudio.field_components.mlp import MLP +from nerfstudio.fields.base_field import Field + + +class TensoRFField(Field): + """TensoRF Field""" + + def __init__( + self, + aabb, + # the aabb bounding box of the dataset + feature_encoding: Encoding = Identity(in_dim=3), + # the encoding method used for appearance encoding outputs + direction_encoding: Encoding = Identity(in_dim=3), + # the encoding method used for ray direction + density_encoding: Encoding = Identity(in_dim=3), + # the tensor encoding method used for scene density + color_encoding: Encoding = Identity(in_dim=3), + # the tensor encoding method used for scene color + appearance_dim: int = 27, + # the number of dimensions for the appearance embedding + head_mlp_num_layers: int = 2, + # number of layers for the MLP + head_mlp_layer_width: int = 128, + # layer width for the MLP + use_sh: bool = False, + # whether to use spherical harmonics as the feature decoding function + sh_levels: int = 2, + # number of levels to use for spherical harmonics + ) -> None: + super().__init__() + self.aabb = Parameter(aabb, requires_grad=False) + self.feature_encoding = feature_encoding + self.direction_encoding = direction_encoding + self.density_encoding = density_encoding + self.color_encoding = color_encoding + + self.mlp_head = MLP( + in_dim=appearance_dim + 3 + self.direction_encoding.get_out_dim() + self.feature_encoding.get_out_dim(), + num_layers=head_mlp_num_layers, + layer_width=head_mlp_layer_width, + activation=nn.ReLU(), + ) + + self.use_sh = use_sh + + if self.use_sh: + self.sh = SHEncoding(sh_levels) + self.B = nn.Linear( + in_features=self.color_encoding.get_out_dim(), out_features=3 * self.sh.get_out_dim(), bias=False + ) + else: + self.B = nn.Linear(in_features=self.color_encoding.get_out_dim(), out_features=appearance_dim, bias=False) + + self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_head.get_out_dim(), activation=nn.Sigmoid()) + + def get_density(self, ray_samples: RaySamples): + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + density = self.density_encoding(positions) + density_enc = torch.sum(density, dim=-1)[:, :, None] + relu = torch.nn.ReLU() + density_enc = relu(density_enc) + return density_enc + + def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None) -> TensorType: + d = ray_samples.frustums.directions + positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + rgb_features = self.color_encoding(positions) + rgb_features = self.B(rgb_features) + + d_encoded = self.direction_encoding(d) + rgb_features_encoded = self.feature_encoding(rgb_features) + + if self.use_sh: + sh_mult = self.sh(d)[:, :, None] + rgb_sh = rgb_features.view(sh_mult.shape[0], sh_mult.shape[1], 3, sh_mult.shape[-1]) + rgb = torch.relu(torch.sum(sh_mult * rgb_sh, dim=-1) + 0.5) + else: + out = self.mlp_head(torch.cat([rgb_features, d, rgb_features_encoded, d_encoded], dim=-1)) # type: ignore + rgb = self.field_output_rgb(out) + + return rgb + + def forward( + self, ray_samples: RaySamples, mask: Optional[TensorType] = None, bg_color: Optional[TensorType] = None + ): + if mask is not None and bg_color is not None: + base_density = torch.zeros(ray_samples.shape)[:, :, None].to(mask.device) + base_rgb = bg_color.repeat(ray_samples[:, :, None].shape) + if mask.any(): + input_rays = ray_samples[mask, :] + density = self.get_density(input_rays) + rgb = self.get_outputs(input_rays, None) + + base_density[mask] = density + base_rgb[mask] = rgb + + base_density.requires_grad_() + base_rgb.requires_grad_() + + density = base_density + rgb = base_rgb + else: + density = self.get_density(ray_samples) + rgb = self.get_outputs(ray_samples, None) + + return {FieldHeadNames.DENSITY: density, FieldHeadNames.RGB: rgb} diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py new file mode 100644 index 000000000..a0c771d2c --- /dev/null +++ b/nerfstudio/models/tensorf.py @@ -0,0 +1,280 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +TensorRF implementation. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Dict, List, Tuple, Type + +import numpy as np +import torch +from torch.nn import Parameter +from torchmetrics import PeakSignalNoiseRatio +from torchmetrics.functional import structural_similarity_index_measure +from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity + +from nerfstudio.cameras.rays import RayBundle +from nerfstudio.configs.config_utils import to_immutable_dict +from nerfstudio.engine.callbacks import ( + TrainingCallback, + TrainingCallbackAttributes, + TrainingCallbackLocation, +) +from nerfstudio.field_components.encodings import NeRFEncoding, TensorVMEncoding +from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.fields.tensorf_field import TensoRFField +from nerfstudio.model_components.losses import MSELoss +from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler +from nerfstudio.model_components.renderers import ( + AccumulationRenderer, + DepthRenderer, + RGBRenderer, +) +from nerfstudio.model_components.scene_colliders import AABBBoxCollider +from nerfstudio.models.base_model import Model, VanillaModelConfig +from nerfstudio.utils import colormaps, colors, misc + + +@dataclass +class TensoRFModelConfig(VanillaModelConfig): + """TensoRF model config""" + + _target: Type = field(default_factory=lambda: TensoRFModel) + """target class to instantiate""" + init_resolution: int = 128 + """initial render resolution""" + final_resolution: int = 300 + """final render resolution""" + upsampling_iters: Tuple[int, ...] = (2000, 3000, 4000, 5500, 7000) + """specifies a list of iteration step numbers to perform upsampling""" + loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss": 1.0}) + """Loss specific weights.""" + num_samples: int = 256 + """Number of samples in field evaluation""" + num_den_components: int = 16 + """Number of components in density encoding""" + num_color_components: int = 48 + """Number of components in color encoding""" + appearance_dim: int = 27 + """Number of channels for color encoding""" + + +class TensoRFModel(Model): + """TensoRF Model + + Args: + config: TensoRF configuration to instantiate model + """ + + def __init__( + self, + config: TensoRFModelConfig, + **kwargs, + ) -> None: + self.init_resolution = config.init_resolution + self.upsampling_iters = config.upsampling_iters + self.num_den_components = config.num_den_components + self.num_color_components = config.num_color_components + self.appearance_dim = config.appearance_dim + self.upsampling_steps = ( + np.round( + np.exp( + np.linspace( + np.log(config.init_resolution), + np.log(config.final_resolution), + len(config.upsampling_iters) + 1, + ) + ) + ) + .astype("int") + .tolist()[1:] + ) + super().__init__(config=config, **kwargs) + + def get_training_callbacks( + self, training_callback_attributes: TrainingCallbackAttributes + ) -> List[TrainingCallback]: + + # the callback that we want to run every X iterations after the training iteration + def reinitialize_optimizer( + self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument + ): + resolution = self.upsampling_steps.pop(0) + + # upsample the position and direction grids + self.field.density_encoding.upsample_grid(resolution) + self.field.color_encoding.upsample_grid(resolution) + + # reinitialize the encodings optimizer + optimizers_config = training_callback_attributes.optimizers.config + enc = training_callback_attributes.pipeline.get_param_groups()["encodings"] + lr_init = optimizers_config["encodings"]["optimizer"].lr + + training_callback_attributes.optimizers.optimizers["encodings"] = optimizers_config["encodings"][ + "optimizer" + ].setup(params=enc) + if optimizers_config["encodings"]["scheduler"]: + training_callback_attributes.optimizers.schedulers["encodings"] = optimizers_config["encodings"][ + "scheduler" + ].setup(optimizer=training_callback_attributes.optimizers.optimizers["encodings"], lr_init=lr_init) + + callbacks = [ + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + iters=self.upsampling_iters, + func=reinitialize_optimizer, + args=[self, training_callback_attributes], + ) + ] + return callbacks + + def populate_modules(self): + """Set the fields and modules""" + 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, + ) + + 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) + + self.field = TensoRFField( + self.scene_box.aabb, + feature_encoding=feature_encoding, + direction_encoding=direction_encoding, + density_encoding=density_encoding, + color_encoding=color_encoding, + appearance_dim=self.appearance_dim, + head_mlp_num_layers=2, + head_mlp_layer_width=128, + use_sh=False, + ) + + # samplers + self.sampler_uniform = UniformSampler(num_samples=self.config.num_samples, single_jitter=True) + self.sampler_pdf = PDFSampler(num_samples=self.config.num_samples // 2, single_jitter=True) + + # renderers + self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) + self.renderer_accumulation = AccumulationRenderer() + self.renderer_depth = DepthRenderer() + + # losses + self.rgb_loss = MSELoss() + + # metrics + self.psnr = PeakSignalNoiseRatio(data_range=1.0) + self.ssim = structural_similarity_index_measure + self.lpips = LearnedPerceptualImagePatchSimilarity() + + # colliders + if self.config.enable_collider: + self.collider = AABBBoxCollider(scene_box=self.scene_box) + + def get_param_groups(self) -> Dict[str, List[Parameter]]: + param_groups = {} + + param_groups["fields"] = ( + list(self.field.mlp_head.parameters()) + + list(self.field.B.parameters()) + + list(self.field.field_output_rgb.parameters()) + ) + param_groups["encodings"] = list(self.field.color_encoding.parameters()) + list( + self.field.density_encoding.parameters() + ) + + return param_groups + + def get_outputs(self, ray_bundle: RayBundle): + # uniform sampling + ray_samples_uniform = self.sampler_uniform(ray_bundle) + dens = self.field.get_density(ray_samples_uniform) + weights = ray_samples_uniform.get_weights(dens) + coarse_accumulation = self.renderer_accumulation(weights) + acc_mask = torch.where(coarse_accumulation < 0.0001, False, True).reshape(-1) + + # pdf sampling + ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights) + + # fine field: + field_outputs_fine = self.field.forward(ray_samples_pdf, acc_mask, colors.WHITE.to(weights.device)) + + weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) + + accumulation = self.renderer_accumulation(weights_fine) + depth = self.renderer_depth(weights_fine, ray_samples_pdf) + + rgb = self.renderer_rgb( + rgb=field_outputs_fine[FieldHeadNames.RGB], + weights=weights_fine, + ) + + rgb = torch.where(accumulation < 0, colors.WHITE.to(rgb.device), rgb) + accumulation = torch.clamp(accumulation, min=0) + + outputs = {"rgb": rgb, "accumulation": accumulation, "depth": depth} + return outputs + + def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: + # Scaling metrics by coefficients to create the losses. + device = outputs["rgb"].device + image = batch["image"].to(device) + + rgb_loss = self.rgb_loss(image, outputs["rgb"]) + + loss_dict = {"rgb_loss": rgb_loss} + loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) + return loss_dict + + def get_image_metrics_and_images( + self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor] + ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]: + image = batch["image"].to(outputs["rgb"].device) + rgb = outputs["rgb"] + acc = colormaps.apply_colormap(outputs["accumulation"]) + depth = colormaps.apply_depth_colormap( + outputs["depth"], + accumulation=outputs["accumulation"], + near_plane=self.config.collider_params["near_plane"], + far_plane=self.config.collider_params["far_plane"], + ) + + combined_rgb = torch.cat([image, rgb], dim=1) + + # Switch images from [H, W, C] to [1, C, H, W] for metrics computations + image = torch.moveaxis(image, -1, 0)[None, ...] + rgb = torch.moveaxis(rgb, -1, 0)[None, ...] + + psnr = self.psnr(image, rgb) + ssim = self.ssim(image, rgb) + lpips = self.lpips(image, rgb) + + metrics_dict = { + "psnr": float(psnr.item()), + "ssim": float(ssim.item()), + "lpips": float(lpips.item()), + } + images_dict = {"img": combined_rgb, "accumulation": acc, "depth": depth} + return metrics_dict, images_dict