From dc1b8092111975e2d349a7fca53dd12966ddc1a5 Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Wed, 2 Aug 2023 09:54:20 -0700 Subject: [PATCH] remove barf --- nerfstudio/configs/method_configs.py | 2 +- nerfstudio/fields/nerfacto_field.py | 51 +--------------------------- nerfstudio/models/nerfacto.py | 13 ------- 3 files changed, 2 insertions(+), 64 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 4a4cc4dc2..f9a5eb1f2 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -121,7 +121,7 @@ method_configs["nerfacto"] = TrainerConfig( }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), - vis="viewer_beta", + vis="viewer", ) method_configs["nerfacto-big"] = TrainerConfig( diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index e528a2164..bb8bd7daf 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -65,8 +65,6 @@ class NerfactoField(Field): use_pred_normals: whether to use predicted normals use_average_appearance_embedding: whether to use average appearance embedding or zeros for inference spatial_distortion: spatial distortion to apply to the scene - bundle_adjust: whether to use bundle adjustment - coarse_to_fine_iters: iterations (percentage of max iters) at which bundle adjustment is active """ aabb: Tensor @@ -92,14 +90,11 @@ class NerfactoField(Field): use_transient_embedding: bool = False, use_semantics: bool = False, num_semantic_classes: int = 100, - freq_warmup: int = 0, pass_semantic_gradients: bool = False, use_pred_normals: bool = False, use_average_appearance_embedding: bool = False, spatial_distortion: Optional[SpatialDistortion] = None, implementation: Literal["tcnn", "torch"] = "tcnn", - bundle_adjust: bool = False, - coarse_to_fine_iters: Optional[Tuple[float, float]] = None, ) -> None: super().__init__() @@ -120,9 +115,6 @@ class NerfactoField(Field): self.use_pred_normals = use_pred_normals self.pass_semantic_gradients = pass_semantic_gradients self.base_res = base_res - self.freq_warmup = freq_warmup - self.bundle_adjust = bundle_adjust - self.coarse_to_fine_iters = coarse_to_fine_iters self.step = 0 self.direction_encoding = SHEncoding( @@ -208,44 +200,6 @@ class NerfactoField(Field): implementation=implementation, ) - def step_cb(self, step): - self.step = step - - def get_freq_mask(self, B, N, D, device): - import numpy as np - - if self.bundle_adjust: - # set masks for hash grid levels - self.max_iters: int = GLOBAL_BUFFER.get("max_iter") - assert self.coarse_to_fine_iters is not None - start, end = self.coarse_to_fine_iters - assert ( - start >= 0 and end >= 0 - ), f"start and end iterations for bundle adjustment have to be positive, got start = {start} and end = {end}" - L = N - # From https://arxiv.org/pdf/2104.06405.pdf equation 14 - alpha = ((self.step / self.max_iters) - start) / (end - start) * L - k = torch.arange(L, dtype=torch.float32, device=device) - mask_vals = (1 - (alpha - k).clamp_(min=0, max=1).mul_(np.pi).cos_()) / 2 - else: - ts = np.linspace(self.step * N, self.step, N) - mask_vals = np.interp(ts, [0, self.freq_warmup], [0, 1]) - mask_vals = torch.from_numpy(mask_vals) - m = mask_vals[None, ..., None].to(device) - return m.repeat((B, 1, D)) - - def mask_freqs(self, hashgrid_outputs): - """ - hashgrid_outputs is a B x N*D array of the hashgrid values - """ - # N is number of levels (usually 16) - # D is feature dim (usually 2) - N = self.mlp_base_grid.num_levels - D = self.mlp_base_grid.features_per_level - hashgrid_outputs = hashgrid_outputs.view(-1, N, D) - masked = self.get_freq_mask(hashgrid_outputs.shape[0], N, D, hashgrid_outputs.device) * hashgrid_outputs - return masked.view(masked.shape[0], -1) - def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: """Computes and returns the densities.""" if self.spatial_distortion is not None: @@ -261,10 +215,7 @@ class NerfactoField(Field): if not self._sample_locations.requires_grad: self._sample_locations.requires_grad = True positions_flat = positions.view(-1, 3) - hashgrid_vecs = self.mlp_base_grid(positions_flat) - hashgrid_vecs = self.mask_freqs(hashgrid_vecs) - h = self.mlp_base_mlp(hashgrid_vecs).view(*ray_samples.frustums.shape, -1) - # h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1) + h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1) density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) self._density_before_activation = density_before_activation diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 8db2a01fc..4d42b7c5a 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -131,11 +131,6 @@ class NerfactoModelConfig(ModelConfig): """Dimension of the appearance embedding.""" camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig(mode="SO3xR3") """Config of the camera optimizer to use""" - bundle_adjust: bool = False - """Whether to bundle adjust (BARF)""" - coarse_to_fine_iters: Optional[Tuple[float, float]] = (0.0, 0.1) - """Iterations (as a percentage of total iterations) at which coarse to fine hash grid optimization starts and ends. - Linear interpolation between (start, end) and full activation of hash grid from end onwards.""" class NerfactoModel(Model): @@ -171,8 +166,6 @@ class NerfactoModel(Model): use_average_appearance_embedding=self.config.use_average_appearance_embedding, appearance_embedding_dim=self.config.appearance_embed_dim, implementation=self.config.implementation, - bundle_adjust=self.config.bundle_adjust, - coarse_to_fine_iters=self.config.coarse_to_fine_iters, ) self.camera_optimizer: CameraOptimizer = self.config.camera_optimizer.setup( @@ -287,12 +280,6 @@ class NerfactoModel(Model): func=self.proposal_sampler.step_cb, ) ) - callbacks.append( - TrainingCallback( - where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - update_every_num_iters=1, - func=self.field.step_cb, - ) ) return callbacks