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"]