From 207959da16f820efabe41d6d950569c341a33d79 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Tue, 21 Feb 2023 17:12:23 -0800 Subject: [PATCH] Warmups --- nerfstudio/configs/method_configs.py | 13 ++++++------ .../datamanagers/dreamfusion_datamanager.py | 11 +++++++--- .../generative/stable_diffusion_utils.py | 2 +- nerfstudio/models/dreamfusion.py | 21 ++++++++++--------- 4 files changed, 27 insertions(+), 20 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 99fbc1c75..e16f2c302 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -349,21 +349,22 @@ method_configs["dreamfusion"] = TrainerConfig( dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, + horizontal_rotation_warmup=4000, ), model=DreamFusionModelConfig( eval_num_rays_per_chunk=1 << 15, distortion_loss_mult=0.02, - interlevel_loss_mult=1.0, + interlevel_loss_mult=100.0, orientation_loss_mult=0.1, max_res=256, sphere_collider=True, initialize_density=False, random_background=True, - proposal_warmup=1000, + proposal_warmup=500, proposal_update_every=2, proposal_weights_anneal_max_num_iters=100, - start_lambertian_training=1000, - start_normals_training=500, + start_lambertian_training=500, + start_normals_training=1000, opacity_loss_mult=0.001, positional_prompting="discrete", guidance_scale=100, @@ -375,8 +376,8 @@ method_configs["dreamfusion"] = TrainerConfig( "scheduler": None, }, "fields": { - "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15), - "scheduler": None, + "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15), + "scheduler": ExponentialDecaySchedulerConfig(warmup_steps=2000, lr_final=1e-6, max_steps=20000), }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), diff --git a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py index 153913d42..9d7af9d3f 100644 --- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py +++ b/nerfstudio/data/datamanagers/dreamfusion_datamanager.py @@ -66,7 +66,7 @@ def random_train_pose( radius_mean: float = 1.0, radius_std: float = 0.1, central_rotation_range: Tuple[float, float] = (0, 360), - vertical_rotation_range: Tuple[float, float] = (-80, 20), + vertical_rotation_range: Tuple[float, float] = (-90, 0), focal_range: Tuple[float, float] = (0.75, 1.35), jitter_std: float = 0.01, center: Tuple[float, float, float] = (0, 0, 0), @@ -169,11 +169,14 @@ class DreamFusionDataManagerConfig(VanillaDataManagerConfig): """Std of radius of camera orbit""" focal_range: Tuple[float, float] = (0.6, 1.2) """Range of focal length""" - vertical_rotation_range: Tuple[float, float] = (-80, 30) + vertical_rotation_range: Tuple[float, float] = (-90, 0) """Range of vertical rotation""" jitter_std: float = 0.05 """Std of camera direction jitter, so we don't just point the cameras towards the center every time""" center: Tuple[float, float, float] = (0, 0, 0) + """Center coordinate of the camera sphere""" + horizontal_rotation_warmup: int = 0 + """How many steps until the full horizontal rotation range is used""" class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-method @@ -244,7 +247,8 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me # ).flatten() # return ray_bundle, {"initialization": False} - # TODO below + horizontal_range = min((step / self.config.horizontal_rotation_warmup), 1) * 180 + cameras, vertical_rotation, central_rotation = random_train_pose( self.config.train_images_per_batch, self.config.train_resolution, @@ -255,6 +259,7 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me vertical_rotation_range=self.config.vertical_rotation_range, jitter_std=self.config.jitter_std, center=self.config.center, + central_rotation_range=(-horizontal_range, horizontal_range), ) ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten() diff --git a/nerfstudio/generative/stable_diffusion_utils.py b/nerfstudio/generative/stable_diffusion_utils.py index d7a94805e..ec5b29544 100644 --- a/nerfstudio/generative/stable_diffusion_utils.py +++ b/nerfstudio/generative/stable_diffusion_utils.py @@ -46,7 +46,7 @@ class PositionalTextEmbeddings: back_prompt: str, front_prompt: str, stable_diffusion: StableDiffusion, - positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete", + positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete" ): self.positional_prompting = positional_prompting self.sd_device = stable_diffusion.device diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 86e45e0e1..230de5f43 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 = 10 - """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 @@ -126,7 +128,7 @@ class DreamFusionModelConfig(ModelConfig): """start training with lambertian shading after this many iterations""" opacity_penalty: bool = True """enables penalty to encourage sparse weights (penalizing for uniform density along ray)""" - opacity_loss_mult: float = 1 + opacity_loss_mult: float = 1e-3 """scale for opacity penalty""" max_res: int = 256 """Maximum resolution of the density field.""" @@ -135,7 +137,7 @@ class DreamFusionModelConfig(ModelConfig): """enables location based prompting""" interpolated_prompting: bool = False """enables interpolated location prompting""" - positional_prompting: Literal["discrete", "interpolated", "base"] = "discrete" + positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete" """ how to incorporate position into prompt""" top_prompt: str = ", overhead view" """appended to prompt for overhead view""" @@ -237,7 +239,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 @@ -267,7 +269,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 @@ -303,6 +305,7 @@ class DreamFusionModel(Model): 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 +355,9 @@ class DreamFusionModel(Model): if self.initialize_density: pos = ray_samples.frustums.get_positions() density_blob = ( - self.config.init_density_strength - * self.density_strength - * torch.exp(-torch.norm(pos, dim=-1) * 10)[..., None] + self.density_strength * self.density_strength * torch.exp(-torch.norm(pos, dim=-1) * 10)[..., None] ) - density += density_blob + density = density + density_blob weights = ray_samples.get_weights(density) weights_list.append(weights)