From 2efd4c88b689334cfa3d18546cc7d77ee1f87998 Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 21 Feb 2023 20:20:48 +0000 Subject: [PATCH] param tweaks, change density bias scheduling --- nerfstudio/configs/method_configs.py | 8 ++++---- nerfstudio/models/dreamfusion.py | 19 ++++++++++--------- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 448b352a8..611d7eadf 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -353,7 +353,7 @@ method_configs["dreamfusion"] = TrainerConfig( model=DreamFusionModelConfig( eval_num_rays_per_chunk=1 << 15, distortion_loss_mult=0.02, - interlevel_loss_mult=100, + interlevel_loss_mult=100.0, orientation_loss_mult=0.1, max_res=256, sphere_collider=True, @@ -371,12 +371,12 @@ method_configs["dreamfusion"] = TrainerConfig( ), optimizers={ "proposal_networks": { - "optimizer": AdamOptimizerConfig(lr=1e-8, eps=1e-15), - "scheduler": WarmupSchedulerConfig(lr_init=1e-8,lr_max=1e-3,lr_final=1e-6,warmup_steps=2000,max_steps=10000), + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None#WarmupSchedulerConfig(lr_init=1e-8,lr_max=1e-3,lr_final=1e-6,warmup_steps=2000,max_steps=20000), }, "fields": { "optimizer": AdamOptimizerConfig(lr=1e-8, eps=1e-15), - "scheduler": WarmupSchedulerConfig(lr_init=1e-8,lr_max=1e-3,lr_final=1e-6,warmup_steps=2000,max_steps=10000), + "scheduler": WarmupSchedulerConfig(lr_init=1e-8,lr_max=1e-3,lr_final=1e-6,warmup_steps=2000,max_steps=20000), }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index fcb76976b..459a09458 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -77,8 +77,10 @@ class DreamFusionModelConfig(ModelConfig): """Randomizes light source per output.""" initialize_density: bool = True """Initialize density in center of scene.""" - init_density_strength: float = 0.1 - """Initial strength of center density""" + taper_range: Tuple[int,int] = (0,1000) + '''Range of step values for the density tapering''' + taper_strength:Tuple[float,float] = (1.0,0.0) + """Strength schedule of center density""" sphere_collider: bool = True """Use spherical collider instead of box""" random_background: bool = True @@ -114,7 +116,7 @@ class DreamFusionModelConfig(ModelConfig): """Slope of the annealing function for the proposal weights.""" proposal_weights_anneal_max_num_iters: int = 500 """Max num iterations for the annealing function.""" - use_single_jitter: bool = False + use_single_jitter: bool = True """Whether use single jitter or not for the proposal networks.""" interlevel_loss_mult: float = 1.0 """Proposal loss multiplier.""" @@ -234,7 +236,7 @@ class DreamFusionModel(Model): single_jitter=self.config.use_single_jitter, update_sched=update_schedule, initial_sampler=UniformSampler(single_jitter=self.config.use_single_jitter), - pdf_histogram_padding=0.001, + pdf_histogram_padding=0.005, ) # renderers @@ -264,7 +266,7 @@ class DreamFusionModel(Model): def taper_density( self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument ): - self.density_strength -= 0.25 + self.density_strength = np.interp(step,self.config.taper_range,self.config.taper_strength) def start_training_normals( self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument @@ -298,8 +300,8 @@ class DreamFusionModel(Model): callbacks = [ TrainingCallback( where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - iters=(150, 300, 450, 600), func=taper_density, + update_every_num_iters=1, args=[self, training_callback_attributes], ), TrainingCallback( @@ -352,11 +354,10 @@ class DreamFusionModel(Model): if self.initialize_density and self.training: pos = ray_samples.frustums.get_positions() density_blob = ( - self.config.init_density_strength - * self.density_strength + self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] ) - density += density_blob + density = density + density_blob accumulation = self.renderer_accumulation(weights) depth = self.renderer_depth(weights, ray_samples)