From 06ec4bd982333d6c3f18038c82f5fdb9bc927a43 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Tue, 14 Feb 2023 14:37:03 -0800 Subject: [PATCH 1/2] Update configs --- nerfstudio/configs/method_configs.py | 7 ++-- nerfstudio/fields/dreamfusion_field.py | 6 ++-- nerfstudio/models/dreamfusion.py | 34 +++++++++++++------- nerfstudio/pipelines/dreamfusion_pipeline.py | 8 ----- 4 files changed, 29 insertions(+), 26 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 7bf71354f..02d4708a8 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -352,7 +352,9 @@ method_configs["dreamfusion"] = TrainerConfig( ), model=DreamFusionModelConfig( eval_num_rays_per_chunk=1 << 15, - distortion_loss_mult=10, + distortion_loss_mult=10.0, + orientation_loss_mult=0.0001, + max_res=256, sphere_collider=True, initialize_density=False, random_background=True, @@ -363,8 +365,7 @@ method_configs["dreamfusion"] = TrainerConfig( start_lambertian_training=1000, start_normals_training=500, ), - alphas_loss_mult=1.0, - opacity_loss_mult=1.0, + interpolated_prompting=False, guidance_scale=100, ), optimizers={ diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py index f3b61e58f..4960fbe48 100644 --- a/nerfstudio/fields/dreamfusion_field.py +++ b/nerfstudio/fields/dreamfusion_field.py @@ -65,6 +65,9 @@ class DreamFusionField(Field): num_layers: int = 2, hidden_dim: int = 64, geo_feat_dim: int = 15, + num_levels: int = 16, + max_res: int = 256, + log2_hashmap_size: int = 19, num_layers_color: int = 3, hidden_dim_color: int = 64, hidden_dim_normal: int = 32, @@ -74,10 +77,7 @@ class DreamFusionField(Field): self.aabb = Parameter(aabb, requires_grad=False) self.geo_feat_dim = geo_feat_dim - num_levels = 16 - max_res = 1024 base_res = 16 - log2_hashmap_size = 19 features_per_level = 2 growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 30454485a..0c7f8c6a9 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -116,6 +116,16 @@ class DreamFusionModelConfig(ModelConfig): """Start training normals after this many iterations""" start_lambertian_training: int = 1000 """start training with lambertian shading after this many iterations""" + alphas_penalty: bool = True + """enables penalty to encourage sparse weights (penalizing for uniform density along ray)""" + alphas_loss_mult: float = 1 + """scale for alphas penalty""" + opacity_penalty: bool = True + """enables penalty to encourage transparent scenes, as in "dreamfields" paper""" + opacity_loss_mult: float = 1 + """scale for opacity penalty""" + max_res: int = 256 + """Maximum resolution of the density field.""" class DreamFusionModel(Model): @@ -145,7 +155,7 @@ class DreamFusionModel(Model): super().populate_modules() # setting up fields - self.field = DreamFusionField(self.scene_box.aabb) + self.field = DreamFusionField(self.scene_box.aabb, max_res=self.config.max_res) # samplers self.density_fns = [] @@ -339,7 +349,7 @@ class DreamFusionModel(Model): shading_weight = 0.0 shaded, shaded_albedo = self.shader_lambertian( - rgb=rgb, normals=pred_normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=True + rgb=rgb, normals=normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=False ) outputs["normals"] = self.shader_normals(normals, weights=accum_mask) @@ -354,11 +364,7 @@ class DreamFusionModel(Model): if samp > 0.5 and not self.training: outputs["train_output"] = outputs["shaded"] elif samp < 0.2 and self.random_background: - rand_bg = torch.zeros(background.shape) - rand_bg.index_fill_(1, torch.tensor([0]), np.random.random_sample()) - rand_bg.index_fill_(1, torch.tensor([1]), np.random.random_sample()) - rand_bg.index_fill_(1, torch.tensor([2]), np.random.random_sample()) - rand_bg = rand_bg.to(self.device) + rand_bg = torch.ones_like(background) * torch.rand(3, device=self.device) outputs["train_output"] = accum_mask * shaded_albedo + rand_bg * accum_mask_inv else: outputs["train_output"] = accum_mask * shaded_albedo + background @@ -377,7 +383,10 @@ class DreamFusionModel(Model): ) assert weights.shape[-1] == 1 - outputs["alphas_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) + if self.config.alphas_penalty: + outputs["alphas_loss"] = ( + torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.alphas_loss_mult + ) return outputs @@ -399,11 +408,12 @@ class DreamFusionModel(Model): loss_dict["orientation_loss"] = 0 loss_dict["pred_normal_loss"] = 0 - loss_dict["alphas_loss"] = outputs["alphas_loss"].mean() + if self.config.alphas_penalty: + loss_dict["alphas_loss"] = self.config.alphas_loss_mult * outputs["alphas_loss"].mean() - loss_dict["opacity_loss"] = -torch.minimum( - (1 - outputs["accumulation"]).mean(), torch.tensor(self.target_transmittance) - ) + if self.config.opacity_penalty: + opacity_loss = -torch.minimum((1 - outputs["accumulation"]).mean(), torch.tensor(self.target_transmittance)) + loss_dict["opacity_loss"] = self.config.opacity_loss_mult * opacity_loss if self.training: loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss( outputs["weights_list"], outputs["ray_samples_list"] diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 3123bfe70..0c40979dd 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -51,14 +51,6 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig): """appended to prompt for front view""" back_prompt: str = ", back view" """appended to prompt for back view""" - alphas_penalty: bool = True - """enables penalty to encourage sparse weights (penalizing for uniform density along ray)""" - alphas_loss_mult: float = 1 - """scale for alphas penalty""" - opacity_penalty: bool = True - """enables penalty to encourage transparent scenes, as in "dreamfields" paper""" - opacity_loss_mult: float = 1 - """scale for opacity penalty""" guidance_scale: float = 100 """guidance scale for sds loss""" From 4db2410b40ecf6a02cfa78cfe3442d0a879f3cf8 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Wed, 15 Feb 2023 13:43:49 -0800 Subject: [PATCH 2/2] Fix opacity loss, use unform sampling first round of prop --- nerfstudio/configs/method_configs.py | 5 ++-- nerfstudio/model_components/ray_samplers.py | 17 ++++++++++-- nerfstudio/models/dreamfusion.py | 29 +++++++++------------ 3 files changed, 30 insertions(+), 21 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 02d4708a8..d0c76262d 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -341,7 +341,7 @@ method_configs["dreamfusion"] = TrainerConfig( method_name="dreamfusion", steps_per_eval_batch=50, steps_per_eval_image=50, - steps_per_save=2000, + steps_per_save=200, max_num_iterations=30000, mixed_precision=True, pipeline=DreamfusionPipelineConfig( @@ -353,7 +353,7 @@ method_configs["dreamfusion"] = TrainerConfig( model=DreamFusionModelConfig( eval_num_rays_per_chunk=1 << 15, distortion_loss_mult=10.0, - orientation_loss_mult=0.0001, + orientation_loss_mult=0.1, max_res=256, sphere_collider=True, initialize_density=False, @@ -364,6 +364,7 @@ method_configs["dreamfusion"] = TrainerConfig( proposal_weights_anneal_max_num_iters=100, start_lambertian_training=1000, start_normals_training=500, + opacity_loss_mult=0.001, ), interpolated_prompting=False, guidance_scale=100, diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index 5df52d2fa..08ee1fd95 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -509,7 +509,16 @@ class VolumetricSampler(Sampler): class ProposalNetworkSampler(Sampler): - """Sampler that uses a proposal network to generate samples.""" + """Sampler that uses a proposal network to generate samples. + + Args: + num_proposal_samples_per_ray: Number of samples to generate per ray for each proposal step. + num_nerf_samples_per_ray: Number of samples to generate per ray for the NERF model. + num_proposal_network_iterations: Number of proposal network iterations to run. + single_jitter: Use a same random jitter for all samples along a ray. + update_sched: A function that takes the iteration number of steps between updates. + initial_sampler: Sampler to use for the first iteration. Uses UniformLinDispPiecewise if not set. + """ def __init__( self, @@ -518,6 +527,7 @@ class ProposalNetworkSampler(Sampler): num_proposal_network_iterations: int = 2, single_jitter: bool = False, update_sched: Callable = lambda x: 1, + initial_sampler: Optional[Sampler] = None, ) -> None: super().__init__() self.num_proposal_samples_per_ray = num_proposal_samples_per_ray @@ -528,7 +538,10 @@ class ProposalNetworkSampler(Sampler): raise ValueError("num_proposal_network_iterations must be >= 1") # samplers - self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter) + if initial_sampler is None: + self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter) + else: + self.initial_sampler = initial_sampler self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter) self._anneal = 1.0 diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 0c7f8c6a9..21484214e 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -41,7 +41,10 @@ from nerfstudio.model_components.losses import ( orientation_loss, pred_normal_loss, ) -from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler +from nerfstudio.model_components.ray_samplers import ( + ProposalNetworkSampler, + UniformSampler, +) from nerfstudio.model_components.renderers import ( AccumulationRenderer, DepthRenderer, @@ -116,12 +119,8 @@ class DreamFusionModelConfig(ModelConfig): """Start training normals after this many iterations""" start_lambertian_training: int = 1000 """start training with lambertian shading after this many iterations""" - alphas_penalty: bool = True - """enables penalty to encourage sparse weights (penalizing for uniform density along ray)""" - alphas_loss_mult: float = 1 - """scale for alphas penalty""" opacity_penalty: bool = True - """enables penalty to encourage transparent scenes, as in "dreamfields" paper""" + """enables penalty to encourage sparse weights (penalizing for uniform density along ray)""" opacity_loss_mult: float = 1 """scale for opacity penalty""" max_res: int = 256 @@ -183,10 +182,9 @@ class DreamFusionModel(Model): num_proposal_network_iterations=self.config.num_proposal_iterations, single_jitter=self.config.use_single_jitter, update_sched=update_schedule, + initial_sampler=UniformSampler(single_jitter=self.config.use_single_jitter), ) - # self.sampler_uniform = UniformSampler(num_samples=self.num_samples, single_jitter=True) - # renderers self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) self.renderer_accumulation = AccumulationRenderer() @@ -373,7 +371,7 @@ class DreamFusionModel(Model): if self.training or True: outputs["rendered_orientation_loss"] = orientation_loss( - weights, field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions + weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions ) outputs["rendered_pred_normal_loss"] = pred_normal_loss( @@ -383,9 +381,9 @@ class DreamFusionModel(Model): ) assert weights.shape[-1] == 1 - if self.config.alphas_penalty: - outputs["alphas_loss"] = ( - torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.alphas_loss_mult + if self.config.opacity_penalty: + outputs["opacity_loss"] = ( + torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult ) return outputs @@ -408,12 +406,9 @@ class DreamFusionModel(Model): loss_dict["orientation_loss"] = 0 loss_dict["pred_normal_loss"] = 0 - if self.config.alphas_penalty: - loss_dict["alphas_loss"] = self.config.alphas_loss_mult * outputs["alphas_loss"].mean() - if self.config.opacity_penalty: - opacity_loss = -torch.minimum((1 - outputs["accumulation"]).mean(), torch.tensor(self.target_transmittance)) - loss_dict["opacity_loss"] = self.config.opacity_loss_mult * opacity_loss + loss_dict["opacity_loss"] = self.config.opacity_loss_mult * outputs["opacity_loss"].mean() + if self.training: loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss( outputs["weights_list"], outputs["ray_samples_list"]