First pass opacity / sparsity losses

nerfstudio/engine/dreamfusion_trainer.py

@@ -60,6 +60,10 @@ class DreamfusionTrainer(Trainer):
         grad = model_outputs["grad"]
         self.grad_scaler.scale(latents).backward(gradient=grad, retain_graph=True)
         normals_loss = loss_dict["orientation_loss"] + loss_dict["pred_normal_loss"]
+        if self.pipeline.config.alphas_penalty:
+            normals_loss += loss_dict["alphas_loss"]
+        if self.pipeline.config.opacity_penalty:
+            normals_loss += loss_dict["opacity_loss"]
         self.grad_scaler.scale(normals_loss).backward()  # type: ignore
         self.optimizers.optimizer_scaler_step_all(self.grad_scaler)
 

nerfstudio/models/dreamfusion.py

@@ -259,6 +259,9 @@ class DreamFusionModel(Model):
             field_outputs[FieldHeadNames.PRED_NORMALS],
         )
 
+        assert weights.shape[-1] == 1
+        outputs["alphas_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01).mean()
+
         return outputs
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
@@ -276,6 +279,9 @@ class DreamFusionModel(Model):
         loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean(
             outputs["rendered_pred_normal_loss"]
         )
+
+        loss_dict["alphas_loss"] = outputs["alphas_loss"]
+
         return loss_dict
 
     def get_image_metrics_and_images(

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -48,10 +48,11 @@ 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)"""
     opacity_penalty: bool = True
-    """enables penalty to encourage transparent scenes"""
-    target_transmittance: float = 0.5
+    """enables penalty to encourage transparent scenes, as in "dreamfields" paper"""
+    target_transmittance: float = 0.8
     """target transmittance for opacity penalty"""
 
     guidance_scale: float = 100
@@ -144,6 +145,9 @@ class DreamfusionPipeline(VanillaPipeline):
         if self.config.opacity_penalty:
             accum_mean = np.mean(1.0 - accumulation)
             sds_loss *= np.min((0.5, accum_mean))
+            loss_dict["opacity_loss"] = -torch.minimum(
+                (1 - model_outputs["accumulation"]).mean(), torch.tensor(self.config.target_transmittance)
+            )
 
         model_outputs["latents"] = latents
         model_outputs["grad"] = grad