diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 5a8c82d16..cf1453a22 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -41,7 +41,7 @@ from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParser from nerfstudio.data.datasets.depth_dataset import DepthDataset from nerfstudio.data.datasets.sdf_dataset import SDFDataset from nerfstudio.data.datasets.semantic_dataset import SemanticDataset -from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig, AdamWOptimizerConfig +from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig from nerfstudio.engine.schedulers import ( CosineDecaySchedulerConfig, ExponentialDecaySchedulerConfig, @@ -629,19 +629,18 @@ method_configs["neuralangelo"] = TrainerConfig( "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), # "scheduler": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=500000), "scheduler": MultiStepWarmupSchedulerConfig(warm_up_end=5000, milestones=[300_000, 400_000], gamma=0.1), - }, "field_background": { # "optimizer": AdamWOptimizerConfig(lr=1e-3, eps=1e-15), "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), "scheduler": MultiStepWarmupSchedulerConfig(warm_up_end=5000, milestones=[300_000, 400_000], gamma=0.1), - }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), vis="viewer", ) + def merge_methods(methods, method_descriptions, new_methods, new_descriptions, overwrite=True): """Merge new methods and descriptions into existing methods and descriptions. Args: diff --git a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py index d000515aa..708676ed8 100644 --- a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py @@ -33,6 +33,7 @@ from nerfstudio.data.dataparsers.base_dataparser import ( DataparserOutputs, ) from nerfstudio.data.scene_box import SceneBox + # from nerfstudio.utils.images import BasicImages from nerfstudio.utils.io import load_from_json @@ -229,7 +230,7 @@ class SDFStudio(DataParser): intrinsics[0, 2] += 200 height, width = 1200, 1600 meta["height"], meta["width"] = height, width - + if self.config.include_mono_prior: assert meta["has_mono_prior"] # load mono depth @@ -347,7 +348,10 @@ class SDFStudio(DataParser): if pairs_path.exists() and split == "train" and self.config.load_pairs: with open(pairs_path, "r") as f: pairs = f.readlines() - split_ext = lambda x: x.split(".")[0] + + def split_ext(x): + return x.split(".")[0] + pairs_srcs = [] for sources_line in pairs: sources_array = [int(split_ext(img_name)) for img_name in sources_line.split(" ")] diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py index 458e035f7..dd8a70480 100644 --- a/nerfstudio/engine/optimizers.py +++ b/nerfstudio/engine/optimizers.py @@ -60,6 +60,7 @@ class AdamOptimizerConfig(OptimizerConfig): weight_decay: float = 0 """The weight decay to use.""" + @dataclass class AdamWOptimizerConfig(OptimizerConfig): """Basic optimizer config with AdamW""" @@ -68,6 +69,7 @@ class AdamWOptimizerConfig(OptimizerConfig): weight_decay: float = 0 """The weight decay to use.""" + @dataclass class RAdamOptimizerConfig(OptimizerConfig): """Basic optimizer config with RAdam""" diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py index 501eacb52..6d8190ccf 100644 --- a/nerfstudio/engine/schedulers.py +++ b/nerfstudio/engine/schedulers.py @@ -17,7 +17,7 @@ from abc import abstractmethod from dataclasses import dataclass, field -from typing import Literal, Optional, Tuple, Type, Any, List +from typing import Literal, Optional, Tuple, Type, List import numpy as np from torch.optim import Optimizer, lr_scheduler @@ -173,6 +173,7 @@ class CosineDecayScheduler(Scheduler): scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=func) return scheduler + @dataclass class MultiStepWarmupSchedulerConfig(SchedulerConfig): """Basic scheduler config with self-defined exponential decay schedule""" @@ -186,6 +187,7 @@ class MultiStepWarmupSchedulerConfig(SchedulerConfig): gamma: float = 0.33 """The learning rate decay factor.""" + class MultiStepWarmupScheduler(Scheduler): """Starts with a flat lr schedule until it reaches N epochs then applies a given scheduler""" @@ -196,8 +198,8 @@ class MultiStepWarmupScheduler(Scheduler): if step < self.config.warm_up_end: learning_factor = step / self.config.warm_up_end else: - index = np.searchsorted(self.config.milestones, step, side='left') - learning_factor = self.config.gamma ** index + index = np.searchsorted(self.config.milestones, step, side="left") + learning_factor = self.config.gamma**index return learning_factor scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=func) diff --git a/nerfstudio/field_components/encodings.py b/nerfstudio/field_components/encodings.py index 1ed873884..c8f536f79 100644 --- a/nerfstudio/field_components/encodings.py +++ b/nerfstudio/field_components/encodings.py @@ -169,10 +169,10 @@ class NeRFEncoding(Encoding): if self.in_dim is None: raise ValueError("Input dimension has not been set") out_dim = self.in_dim * self.num_frequencies * 2 - + if self.off_axis: out_dim = self.P.shape[1] * self.num_frequencies * 2 - + if self.include_input: out_dim += self.in_dim return out_dim @@ -776,7 +776,6 @@ class PeriodicVolumeEncoding(Encoding): hash_init_scale: float = 0.001, smoothstep: bool = False, ) -> None: - super().__init__(in_dim=3) self.num_levels = num_levels self.features_per_level = features_per_level diff --git a/nerfstudio/fields/sdf_field.py b/nerfstudio/fields/sdf_field.py index 6d8e0a274..a8491ea12 100644 --- a/nerfstudio/fields/sdf_field.py +++ b/nerfstudio/fields/sdf_field.py @@ -69,6 +69,7 @@ except ModuleNotFoundError: # beta = self.beta.abs() + self.beta_min # return beta + class LearnedVariance(nn.Module): """Variance network in NeuS @@ -197,10 +198,10 @@ class SDFField(Field): self.divide_factor = self.config.divide_factor self.num_levels = self.config.num_levels - self.max_res = self.config.max_res - self.base_res = self.config.base_res - self.log2_hashmap_size = self.config.log2_hashmap_size - self.features_per_level = self.config.features_per_level + self.max_res = self.config.max_res + self.base_res = self.config.base_res + self.log2_hashmap_size = self.config.log2_hashmap_size + self.features_per_level = self.config.features_per_level self.smoothstep = self.config.smoothstep self.growth_factor = np.exp((np.log(config.max_res) - np.log(config.base_res)) / (config.num_levels - 1)) @@ -236,7 +237,6 @@ class SDFField(Field): print("using tensor vm") self.encoding = TensorVMEncoding(128, 24, smoothstep=config.smoothstep) - # we concat inputs position ourselves self.position_encoding = NeRFEncoding( in_dim=3, @@ -260,7 +260,6 @@ class SDFField(Field): # deviation_network to compute alpha from sdf from NeuS self.deviation_network = LearnedVariance(init_val=self.config.beta_init) - # diffuse and specular tint layer if self.config.use_diffuse_color: self.diffuse_color_pred = nn.Linear(self.config.geo_feat_dim, 3) @@ -290,15 +289,15 @@ class SDFField(Field): dims = [in_dim] + dims + [3] self.num_layers_color = len(dims) - for l in range(0, self.num_layers_color - 1): - out_dim = dims[l + 1] - lin = nn.Linear(dims[l], out_dim) + for i in range(0, self.num_layers_color - 1): + out_dim = dims[i + 1] + lin = nn.Linear(dims[i], out_dim) torch.nn.init.kaiming_uniform_(lin.weight.data) torch.nn.init.zeros_(lin.bias.data) if self.config.weight_norm: lin = nn.utils.weight_norm(lin) - setattr(self, "clin" + str(l), lin) + setattr(self, "clin" + str(i), lin) self.softplus = nn.Softplus(beta=100) self.relu = nn.ReLU() @@ -356,7 +355,7 @@ class SDFField(Field): def update_mask(self, level: int): self.hash_encoding_mask[:] = 1.0 - self.hash_encoding_mask[level * self.features_per_level:] = 0 + self.hash_encoding_mask[level * self.features_per_level :] = 0 def forward_geonetwork(self, inputs: Float[Tensor, "*batch 3"]) -> Float[Tensor, "*batch geo_features+1"]: """forward the geonetwork""" @@ -398,7 +397,7 @@ class SDFField(Field): hidden_output = self.forward_geonetwork(positions_flat).view(*ray_samples.frustums.shape, -1) sdf, _ = torch.split(hidden_output, [1, self.config.geo_feat_dim], dim=-1) return sdf - + def set_numerical_gradients_delta(self, delta: float) -> None: """Set the delta value for numerical gradient.""" self.numerical_gradients_delta = delta @@ -573,7 +572,7 @@ class SDFField(Field): inputs = self.spatial_distortion(inputs) points_norm = inputs.norm(dim=-1) # compute gradient in constracted space - + inputs.requires_grad_(True) with torch.enable_grad(): hidden_output = self.forward_geonetwork(inputs) @@ -584,11 +583,18 @@ class SDFField(Field): skip_spatial_distortion=True, return_sdf=True, ) - sampled_sdf = sampled_sdf.view(-1, *ray_samples.frustums.directions.shape[:-1]).permute(1, 2, 0).contiguous() + sampled_sdf = ( + sampled_sdf.view(-1, *ray_samples.frustums.directions.shape[:-1]).permute(1, 2, 0).contiguous() + ) else: d_output = torch.ones_like(sdf, requires_grad=False, device=sdf.device) gradients = torch.autograd.grad( - outputs=sdf, inputs=inputs, grad_outputs=d_output, create_graph=True, retain_graph=True, only_inputs=True + outputs=sdf, + inputs=inputs, + grad_outputs=d_output, + create_graph=True, + retain_graph=True, + only_inputs=True, )[0] sampled_sdf = None diff --git a/nerfstudio/models/neuralangelo.py b/nerfstudio/models/neuralangelo.py index 5039cb1d1..88f796b6b 100644 --- a/nerfstudio/models/neuralangelo.py +++ b/nerfstudio/models/neuralangelo.py @@ -19,24 +19,17 @@ Implementation of Neuralangelo model. from __future__ import annotations from dataclasses import dataclass, field -from typing import Dict, List, Tuple, Type +from typing import Dict, List, Type -import numpy as np import torch -from torch.nn import Parameter -from nerfstudio.cameras.rays import RayBundle from nerfstudio.engine.callbacks import ( TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation, ) from nerfstudio.field_components.field_heads import FieldHeadNames -from nerfstudio.fields.density_fields import HashMLPDensityField -from nerfstudio.model_components.losses import interlevel_loss -from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler from nerfstudio.models.neus import NeuSModel, NeuSModelConfig -from nerfstudio.utils import colormaps @dataclass @@ -44,7 +37,7 @@ class NeuralangeloModelConfig(NeuSModelConfig): """Neuralangelo Model Config""" _target: Type = field(default_factory=lambda: NeuralangeloModel) - #TODO move to base model config since it can be used in all models + # TODO move to base model config since it can be used in all models enable_progressive_hash_encoding: bool = True """whether to use progressive hash encoding""" enable_numerical_gradients_schedule: bool = True @@ -60,6 +53,7 @@ class NeuralangeloModelConfig(NeuSModelConfig): steps_per_level: int = 5000 """steps per level of multi-resolution hash encoding""" + class NeuralangeloModel(NeuSModel): """Neuralangelo model @@ -82,24 +76,25 @@ class NeuralangeloModel(NeuSModel): # read the hash encoding parameters from field level_init = self.config.level_init # schedule the delta in numerical gradients computation - num_levels = self.field.num_levels max_res = self.field.max_res base_res = self.field.base_res growth_factor = self.field.growth_factor - + steps_per_level = self.config.steps_per_level - - - init_delta = 1. / base_res - end_delta = 1. / max_res - + + init_delta = 1.0 / base_res + 1.0 / max_res + # compute the delta based on level if self.config.enable_numerical_gradients_schedule: + def set_delta(step): - delta = 1. / (base_res * growth_factor ** ( step / steps_per_level)) - delta = max(1. / max_res, delta) - self.field.set_numerical_gradients_delta(delta * 2.) # TODO because we divide 4 to normalize points to [0, 1] - + delta = 1.0 / (base_res * growth_factor ** (step / steps_per_level)) + delta = max(1.0 / max_res, delta) + self.field.set_numerical_gradients_delta( + delta * 2.0 + ) # TODO because we divide 4 to normalize points to [0, 1] + callbacks.append( TrainingCallback( where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], @@ -107,15 +102,16 @@ class NeuralangeloModel(NeuSModel): func=set_delta, ) ) - + # schedule the current level of multi-resolution hash encoding if self.config.enable_progressive_hash_encoding: + def set_mask(step): - #TODO make this consistent with delta schedule + # TODO make this consistent with delta schedule level = int(step / steps_per_level) + 1 level = max(level, level_init) self.field.update_mask(level) - + callbacks.append( TrainingCallback( where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], @@ -126,16 +122,19 @@ class NeuralangeloModel(NeuSModel): # schedule the curvature loss weight # linear warmup for 5000 steps to 5e-4 and then decay as delta if self.config.enable_curvature_loss_schedule: + def set_curvature_loss_mult_factor(step): if step < self.config.curvature_loss_warmup_steps: factor = step / self.config.curvature_loss_warmup_steps else: - delta = 1. / (base_res * growth_factor ** ( (step - self.config.curvature_loss_warmup_steps) / steps_per_level)) - delta = max(1. / max_res, delta) + delta = 1.0 / ( + base_res * growth_factor ** ((step - self.config.curvature_loss_warmup_steps) / steps_per_level) + ) + delta = max(1.0 / max_res, delta) factor = delta / init_delta self.curvature_loss_multi_factor = factor - + callbacks.append( TrainingCallback( where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], @@ -143,38 +142,38 @@ class NeuralangeloModel(NeuSModel): func=set_curvature_loss_mult_factor, ) ) - - - #TODO switch to analytic gradients after delta is small enough? - + + # TODO switch to analytic gradients after delta is small enough? + return callbacks - + def get_metrics_dict(self, outputs, batch) -> Dict: metrics_dict = super().get_metrics_dict(outputs, batch) - + if self.training: # training statics metrics_dict["activated_encoding"] = self.field.hash_encoding_mask.mean().item() metrics_dict["numerical_gradients_delta"] = self.field.numerical_gradients_delta metrics_dict["curvature_loss_multi"] = self.curvature_loss_multi_factor * self.config.curvature_loss_multi - + return metrics_dict - + def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict: loss_dict = super().get_loss_dict(outputs, batch, metrics_dict) # curvature loss if self.training and self.config.curvature_loss_multi > 0.0: delta = self.field.numerical_gradients_delta - centered_sdf = outputs['field_outputs'][FieldHeadNames.SDF] - sourounding_sdf = outputs['field_outputs']["sampled_sdf"] - + centered_sdf = outputs["field_outputs"][FieldHeadNames.SDF] + sourounding_sdf = outputs["field_outputs"]["sampled_sdf"] + sourounding_sdf = sourounding_sdf.reshape(centered_sdf.shape[:2] + (3, 2)) - + # (a - b)/d - (b -c)/d = (a + c - 2b)/d # ((a - b)/d - (b -c)/d)/d = (a + c - 2b)/(d*d) curvature = (sourounding_sdf.sum(dim=-1) - 2 * centered_sdf) / (delta * delta) - loss_dict["curvature_loss"] = torch.abs(curvature).mean() * self.config.curvature_loss_multi * self.curvature_loss_multi_factor - - return loss_dict + loss_dict["curvature_loss"] = ( + torch.abs(curvature).mean() * self.config.curvature_loss_multi * self.curvature_loss_multi_factor + ) + return loss_dict diff --git a/nerfstudio/models/neus_facto.py b/nerfstudio/models/neus_facto.py index 33e73fd4b..53380576c 100644 --- a/nerfstudio/models/neus_facto.py +++ b/nerfstudio/models/neus_facto.py @@ -219,7 +219,6 @@ class NeuSFactoModel(NeuSModel): # read the hash encoding parameters from field level_init = self.config.level_init # schedule the delta in numerical gradients computation - num_levels = self.field.num_levels max_res = self.field.max_res base_res = self.field.base_res growth_factor = self.field.growth_factor @@ -227,7 +226,7 @@ class NeuSFactoModel(NeuSModel): steps_per_level = self.config.steps_per_level init_delta = 1.0 / base_res - end_delta = 1.0 / max_res + 1.0 / max_res if self.config.enable_numerical_gradients_schedule: @@ -245,8 +244,8 @@ class NeuSFactoModel(NeuSModel): func=set_delta, ) ) - - # schedule the current level of multi-resolution hash encoding + + # schedule the current level of multi-resolution hash encoding if self.config.enable_progressive_hash_encoding: def set_mask(step): @@ -339,7 +338,7 @@ class NeuSFactoModel(NeuSModel): torch.abs(curvature).mean() * self.config.curvature_loss_multi * self.curvature_loss_multi_factor ) return loss_dict - + def get_metrics_dict(self, outputs, batch) -> Dict: metrics_dict = super().get_metrics_dict(outputs, batch) @@ -351,7 +350,6 @@ class NeuSFactoModel(NeuSModel): return metrics_dict - def get_image_metrics_and_images( self, outputs: Dict[str, Any], batch: Dict[str, Any] ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]: