diff --git a/HF_TOKEN b/HF_TOKEN new file mode 100644 index 000000000..82d1555a8 --- /dev/null +++ b/HF_TOKEN @@ -0,0 +1 @@ +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 \ No newline at end of file diff --git a/nerf_accumulation.jpg b/nerf_accumulation.jpg new file mode 100644 index 000000000..d4f10cb84 Binary files /dev/null and b/nerf_accumulation.jpg differ diff --git a/nerf_output.png b/nerf_output.png new file mode 100644 index 000000000..494994a2f Binary files /dev/null and b/nerf_output.png differ diff --git a/nerfstudio/engine/dreamfusion_trainer.py b/nerfstudio/engine/dreamfusion_trainer.py index de456152d..4aa9f1644 100644 --- a/nerfstudio/engine/dreamfusion_trainer.py +++ b/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() diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 38211bacd..e5f7cd341 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/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 diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index b8477ab3c..288dbc6fc 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/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