Small patches and configs

HF_TOKEN

@@ -0,0 +1 @@
+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

nerfstudio/engine/dreamfusion_trainer.py

@@ -34,12 +34,16 @@ class DreamfusionTrainerConfig(TrainerConfig):
     _target: Type = field(default_factory=lambda: DreamfusionTrainer)
     """target class to instantiate"""
     pipeline: DreamfusionPipelineConfig = DreamfusionPipelineConfig()
+    """specifies the pipeline config"""
+    gradient_accumulation_steps: int = 4
+    """number of gradient accumulation steps before optimizer step"""
 
 
 class DreamfusionTrainer(Trainer):
     """Dreamfusion trainer"""
 
     pipeline: DreamfusionPipeline
+    config: DreamfusionTrainerConfig
 
     def __init__(self, config: DreamfusionTrainerConfig, local_rank: int = 0, world_size: int = 1):
         assert isinstance(config, DreamfusionTrainerConfig)
@@ -54,17 +58,26 @@ class DreamfusionTrainer(Trainer):
         Args:
             step: current iteration step to update sampler if using DDP (distributed)
         """
-        self.optimizers.zero_grad_all()
-        model_outputs, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step)
-        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"]
-        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
+        for i in range(self.config.gradient_accumulation_steps):
+            self.optimizers.zero_grad_all()
+            model_outputs, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step)
+            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"]
+            if self.pipeline.config.alphas_penalty:
+                normals_loss += loss_dict["alphas_loss"] * self.config.pipeline.alphas_scale
+            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
+            # 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
+            torch.cuda.empty_cache()
+
         self.optimizers.optimizer_scaler_step_all(self.grad_scaler)
 
         self.grad_scaler.update()

nerfstudio/models/dreamfusion.py

@@ -76,6 +76,14 @@ class DreamFusionModelConfig(ModelConfig):
     """Initial strength of center density"""
     sphere_collider: bool = True
     """Use spherical collider instead of box"""
+    target_transmittance_start: float = 0.8
+    """target transmittance for opacity penalty. This is the percent of the scene that is 
+    background when rendered at the start of training"""
+    target_transmittance_end: float = 0.4
+    """target transmittance for opacity penalty. This is the percent of the scene that is
+    background when rendered at the end of training"""
+    transmittance_end_schedule: int = 1000
+    """number of iterations to reach target_transmittance_end"""
 
 
 class DreamFusionModel(Model):
@@ -97,6 +105,7 @@ class DreamFusionModel(Model):
         self.train_normals = False
         self.train_shaded = False
         self.density_strength = 1.0
+        self.target_transmittance = config.target_transmittance_start
         super().__init__(config=config, **kwargs)
 
     def get_training_callbacks(
@@ -119,6 +128,13 @@ class DreamFusionModel(Model):
         ):
             self.train_shaded = True
 
+        def update_target_transmittance(
+            self, training_callback_attributes: TrainingCallbackAttributes, step: int  # pylint: disable=unused-argument
+        ):
+            self.target_transmittance -= (1 / self.config.transmittance_end_schedule) * (
+                self.config.target_transmittance_start - self.config.target_transmittance_end
+            )
+
         callbacks = [
             TrainingCallback(
                 where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
@@ -138,6 +154,12 @@ class DreamFusionModel(Model):
                 func=start_shaded_training,
                 args=[self, training_callback_attributes],
             ),
+            TrainingCallback(
+                where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
+                iters=(self.config.transmittance_end_schedule,),
+                func=update_target_transmittance,
+                args=[self, training_callback_attributes],
+            ),
         ]
         return callbacks
 
@@ -260,7 +282,7 @@ class DreamFusionModel(Model):
         )
 
         assert weights.shape[-1] == 1
-        outputs["alphas_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01).mean()
+        outputs["alphas_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01)
 
         return outputs
 
@@ -280,7 +302,15 @@ class DreamFusionModel(Model):
             outputs["rendered_pred_normal_loss"]
         )
 
-        loss_dict["alphas_loss"] = outputs["alphas_loss"]
+        loss_dict["alphas_loss"] = outputs["alphas_loss"].mean()
+
+        assert (
+            self.target_transmittance >= self.config.target_transmittance_end
+            and self.target_transmittance <= self.config.target_transmittance_start
+        )
+        loss_dict["opacity_loss"] = -torch.minimum(
+            (1 - outputs["accumulation"]).mean(), torch.tensor(self.target_transmittance)
+        )
 
         return loss_dict
 

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -50,11 +50,12 @@ 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
+    """scale for alphas penalty"""
     opacity_penalty: bool = True
     """enables penalty to encourage transparent scenes, as in "dreamfields" paper"""
-    target_transmittance: float = 0.8
-    """target transmittance for opacity penalty"""
-
+    opacity_scale: float = 1
+    """scale for opacity penalty"""
     guidance_scale: float = 100
     """guidance scale for sds loss"""
 
@@ -145,9 +146,6 @@ 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