Add distortion loss

nerfstudio/configs/method_configs.py

@@ -349,7 +349,10 @@ method_configs["dreamfusion"] = DreamfusionTrainerConfig(
             train_num_rays_per_batch=4096,
             eval_num_rays_per_batch=4096,
         ),
-        model=DreamFusionModelConfig(eval_num_rays_per_chunk=1 << 15),
+        model=DreamFusionModelConfig(eval_num_rays_per_chunk=1 << 15, distortion_loss_mult=10),
+        alphas_loss_mult=1.0,
+        opacity_loss_mult=1.0,
+        guidance_scale=100,
     ),
     optimizers={
         "proposal_networks": {

nerfstudio/engine/dreamfusion_trainer.py

@@ -64,18 +64,20 @@ class DreamfusionTrainer(Trainer):
             latents = model_outputs["latents"]
             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"]
+            loss = loss_dict["orientation_loss"] + loss_dict["pred_normal_loss"]
             if self.pipeline.config.alphas_penalty:
-                normals_loss += loss_dict["alphas_loss"] * self.config.pipeline.alphas_scale
+                loss += loss_dict["alphas_loss"] * self.config.pipeline.alphas_loss_mult
             if self.pipeline.config.opacity_penalty:
-                normals_loss += loss_dict["opacity_loss"] * self.config.pipeline.opacity_scale
-            self.grad_scaler.scale(normals_loss).backward()  # type: ignore
+                loss += loss_dict["opacity_loss"] * self.config.pipeline.opacity_loss_mult
+            loss += loss_dict["distortion_loss"]
+            loss += loss_dict["interlevel_loss"]
+            self.grad_scaler.scale(loss).backward()  # type: ignore
             # I noticed that without these del statements and the empty cache, I was getting OOM errors on the device
             # running stable diffusion, since these variables aren't deleted (normally they would be freed when we
             # exit the frame since they no longer are being referenced, but now we loop back instead of exiting the
             # frame)
             if i != self.config.gradient_accumulation_steps - 1:
-                del normals_loss, latents, grad, model_outputs, loss_dict, metrics_dict
+                del loss, latents, grad, model_outputs, loss_dict, metrics_dict
             torch.cuda.empty_cache()
 
         self.optimizers.optimizer_scaler_step_all(self.grad_scaler)
@@ -83,4 +85,4 @@ class DreamfusionTrainer(Trainer):
         self.grad_scaler.update()
         self.optimizers.scheduler_step_all(step)
 
-        return torch.tensor(0.0), loss_dict, metrics_dict
+        return loss, loss_dict, metrics_dict

nerfstudio/models/dreamfusion.py

@@ -36,6 +36,7 @@ from nerfstudio.fields.density_fields import HashMLPDensityField
 from nerfstudio.fields.dreamfusion_field import DreamFusionField
 from nerfstudio.model_components.losses import (
     MSELoss,
+    distortion_loss,
     interlevel_loss,
     orientation_loss,
     pred_normal_loss,
@@ -106,12 +107,14 @@ class DreamFusionModelConfig(ModelConfig):
     """Arguments for the proposal density fields."""
     proposal_weights_anneal_slope: float = 10.0
     """Slope of the annealing function for the proposal weights."""
-    proposal_weights_anneal_max_num_iters: int = 1000
+    proposal_weights_anneal_max_num_iters: int = 500
     """Max num iterations for the annealing function."""
     use_single_jitter: bool = True
     """Whether use single jitter or not for the proposal networks."""
     interlevel_loss_mult: float = 1.0
     """Proposal loss multiplier."""
+    distortion_loss_mult: float = 1.0
+    """Distortion loss multiplier."""
 
 
 class DreamFusionModel(Model):
@@ -393,6 +396,9 @@ class DreamFusionModel(Model):
             (1 - outputs["accumulation"]).mean(), torch.tensor(self.target_transmittance)
         )
         if self.training:
+            loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss(
+                outputs["weights_list"], outputs["ray_samples_list"]
+            )
             loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
                 outputs["weights_list"], outputs["ray_samples_list"]
             )

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -52,11 +52,11 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig):
     """appended to prompt for back view"""
     alphas_penalty: bool = True
     """enables penalty to encourage sparse weights (penalizing for uniform density along ray)"""
-    alphas_scale: float = 1
+    alphas_loss_mult: float = 1
     """scale for alphas penalty"""
     opacity_penalty: bool = True
     """enables penalty to encourage transparent scenes, as in "dreamfields" paper"""
-    opacity_scale: float = 1
+    opacity_loss_mult: float = 1
     """scale for opacity penalty"""
     guidance_scale: float = 100
     """guidance scale for sds loss"""