diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index dbeb1fc86..f4a95f027 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -58,7 +58,7 @@ method_configs["nerfacto"] = Config( train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, camera_optimizer=CameraOptimizerConfig( - mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3) + mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) ), ), model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), diff --git a/nerfstudio/fields/density_fields.py b/nerfstudio/fields/density_fields.py index af7b534ed..1584deedf 100644 --- a/nerfstudio/fields/density_fields.py +++ b/nerfstudio/fields/density_fields.py @@ -62,11 +62,9 @@ class HashMLPDensityField(Field): features_per_level=2, ) -> None: super().__init__() - self.aabb = Parameter(aabb, requires_grad=False) self.spatial_distortion = spatial_distortion self.use_linear = use_linear - growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) config = { @@ -105,12 +103,6 @@ class HashMLPDensityField(Field): else: positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) positions_flat = positions.view(-1, 3) - # assert all positions are in the range [0, 1] - # otherwise print min and max values - # assert torch.all(positions >= 0.0) and torch.all( - # positions <= 1.0 - # ), f"positions: {positions.min()} {positions.max()}" - if not self.use_linear: density_before_activation = ( self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1).to(positions) diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index f744e66ae..a65bd2128 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -505,12 +505,13 @@ class ProposalNetworkSampler(Sampler): num_nerf_samples_per_ray: int = 32, num_proposal_network_iterations: int = 2, single_jitter: bool = False, + update_sched: Callable = lambda x: 1, ) -> None: super().__init__() self.num_proposal_samples_per_ray = num_proposal_samples_per_ray self.num_nerf_samples_per_ray = num_nerf_samples_per_ray self.num_proposal_network_iterations = num_proposal_network_iterations - + self.update_sched = update_sched if self.num_proposal_network_iterations < 1: raise ValueError("num_proposal_network_iterations must be >= 1") @@ -519,11 +520,18 @@ class ProposalNetworkSampler(Sampler): self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter) self._anneal = 1.0 + self._steps_since_update = 0 + self._step = 0 def set_anneal(self, anneal: float) -> None: """Set the anneal value for the proposal network.""" self._anneal = anneal + def step_cb(self, step): + """Callback to register a training step has passed. This is used to keep track of the sampling schedule""" + self._step = step + self._steps_since_update += 1 + def generate_ray_samples( self, ray_bundle: Optional[RayBundle] = None, @@ -538,6 +546,7 @@ class ProposalNetworkSampler(Sampler): n = self.num_proposal_network_iterations weights = None ray_samples = None + updated = self._steps_since_update > self.update_sched(self._step) or self._step < 10 for i_level in range(n + 1): is_prop = i_level < n num_samples = self.num_proposal_samples_per_ray[i_level] if is_prop else self.num_nerf_samples_per_ray @@ -551,10 +560,17 @@ class ProposalNetworkSampler(Sampler): annealed_weights = torch.pow(weights, self._anneal) ray_samples = self.pdf_sampler(ray_bundle, ray_samples, annealed_weights, num_samples=num_samples) if is_prop: - density = density_fns[i_level](ray_samples.frustums.get_positions()) + if updated: + # always update on the first step or the inf check in grad scaling crashes + density = density_fns[i_level](ray_samples.frustums.get_positions()) + else: + with torch.no_grad(): + density = density_fns[i_level](ray_samples.frustums.get_positions()) weights = ray_samples.get_weights(density) weights_list.append(weights) # (num_rays, num_samples) ray_samples_list.append(ray_samples) + if updated: + self._steps_since_update = 0 assert ray_samples is not None return ray_samples, weights_list, ray_samples_list diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index a45fafaab..6f6293f60 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -61,18 +61,22 @@ class NerfactoModelConfig(ModelConfig): """How far along the ray to stop sampling.""" background_color: Literal["background", "last_sample"] = "last_sample" """Whether to randomize the background color.""" - num_proposal_samples_per_ray: Tuple[int] = (200, 64) + num_proposal_samples_per_ray: Tuple[int] = (256, 96) """Number of samples per ray for the proposal network.""" - num_nerf_samples_per_ray: int = 64 + num_nerf_samples_per_ray: int = 48 """Number of samples per ray for the nerf network.""" + proposal_update_every: int = 5 + """Sample every n steps after the warmup""" + proposal_warmup: int = 5000 + """Scales n from 1 to proposal_update_every over this many steps""" num_proposal_iterations: int = 2 """Number of proposal network iterations.""" use_same_proposal_network: bool = False """Use the same proposal network. Otherwise use different ones.""" proposal_net_args_list: List[Dict] = field( default_factory=lambda: [ - {"hidden_dim": 16, "log2_hashmap_size": 16, "num_levels": 5, "max_res": 64}, - {"hidden_dim": 16, "log2_hashmap_size": 16, "num_levels": 5, "max_res": 256}, + {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 64}, + {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256}, ] ) """Arguments for the proposal density fields.""" @@ -129,17 +133,25 @@ class NerfactoModel(Model): for i in range(num_prop_nets): prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)] network = HashMLPDensityField( - self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args + self.scene_box.aabb, + spatial_distortion=scene_contraction, + **prop_net_args, ) self.proposal_networks.append(network) self.density_fns.extend([network.density_fn for network in self.proposal_networks]) # Samplers + update_schedule = lambda step: np.clip( + np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]), + 1, + self.config.proposal_update_every, + ) self.proposal_sampler = ProposalNetworkSampler( num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray, num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray, num_proposal_network_iterations=self.config.num_proposal_iterations, single_jitter=self.config.use_single_jitter, + update_sched=update_schedule, ) # Collider @@ -186,6 +198,13 @@ class NerfactoModel(Model): func=set_anneal, ) ) + callbacks.append( + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + update_every_num_iters=1, + func=self.proposal_sampler.step_cb, + ) + ) return callbacks def get_outputs(self, ray_bundle: RayBundle):