diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index e16f2c302..a03b8bcf1 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -361,7 +361,7 @@ method_configs["dreamfusion"] = TrainerConfig( initialize_density=False, random_background=True, proposal_warmup=500, - proposal_update_every=2, + proposal_update_every=0, proposal_weights_anneal_max_num_iters=100, start_lambertian_training=500, start_normals_training=1000, @@ -372,7 +372,7 @@ method_configs["dreamfusion"] = TrainerConfig( ), optimizers={ "proposal_networks": { - "optimizer": AdamOptimizerConfig(lr=5e-2, eps=1e-15), + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), "scheduler": None, }, "fields": { diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 6d94ed01d..4dd4b6b1f 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -128,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 = 1e-3 + opacity_loss_mult: float = 1 """scale for opacity penalty""" max_res: int = 256 """Maximum resolution of the density field.""" @@ -303,7 +303,6 @@ 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], @@ -354,9 +353,7 @@ class DreamFusionModel(Model): if self.initialize_density: pos = ray_samples.frustums.get_positions() - density_blob = ( - self.density_strength * self.density_strength * torch.exp(-torch.norm(pos, dim=-1) * 10)[..., None] - ) + density_blob = self.density_strength * torch.exp(-torch.norm(pos, dim=-1) / (2 * 0.04))[..., None] density = density + density_blob weights = ray_samples.get_weights(density) @@ -426,22 +423,19 @@ class DreamFusionModel(Model): else: outputs["train_output"] = outputs["rgb"] - if self.training: - outputs["rendered_orientation_loss"] = orientation_loss( - weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions - ) + outputs["rendered_orientation_loss"] = orientation_loss( + weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions + ) - outputs["rendered_pred_normal_loss"] = pred_normal_loss( - weights.detach(), - field_outputs[FieldHeadNames.NORMALS].detach(), - field_outputs[FieldHeadNames.PRED_NORMALS], - ) + outputs["rendered_pred_normal_loss"] = pred_normal_loss( + weights.detach(), + field_outputs[FieldHeadNames.NORMALS].detach(), + field_outputs[FieldHeadNames.PRED_NORMALS], + ) - assert weights.shape[-1] == 1 - if self.config.opacity_penalty: - outputs["opacity_loss"] = ( - torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult - ) + assert weights.shape[-1] == 1 + 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