From 6c9b37bff9afb03ebb77a6bca409b37843f4f389 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Mon, 6 Feb 2023 17:14:57 -0800 Subject: [PATCH] Add distortion loss --- nerfstudio/configs/method_configs.py | 5 ++++- nerfstudio/engine/dreamfusion_trainer.py | 14 ++++++++------ nerfstudio/models/dreamfusion.py | 8 +++++++- nerfstudio/pipelines/dreamfusion_pipeline.py | 4 ++-- 4 files changed, 21 insertions(+), 10 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 280621282..228c13a62 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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": { diff --git a/nerfstudio/engine/dreamfusion_trainer.py b/nerfstudio/engine/dreamfusion_trainer.py index 66368f350..b0eb163c1 100644 --- a/nerfstudio/engine/dreamfusion_trainer.py +++ b/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 diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index b5c0a3c2d..0208b4eca 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/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"] ) diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index e4d1f3def..d2c03eb56 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/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"""