From 76e5bd9b46004096adc51ab7ef017e2e75f0d462 Mon Sep 17 00:00:00 2001 From: David McAllister Date: Mon, 13 Feb 2023 22:15:13 -0800 Subject: [PATCH] Remove Dreamfusion Trainer (#1423) * Dreamfusion trainer no longer needed * Added typing for grad scaler --- nerfstudio/configs/method_configs.py | 3 +- nerfstudio/engine/dreamfusion_trainer.py | 88 -------------------- nerfstudio/engine/trainer.py | 2 +- nerfstudio/pipelines/base_pipeline.py | 2 + nerfstudio/pipelines/dreamfusion_pipeline.py | 23 ++++- nerfstudio/pipelines/dynamic_batch.py | 4 +- 6 files changed, 28 insertions(+), 94 deletions(-) delete mode 100644 nerfstudio/engine/dreamfusion_trainer.py diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 1a5cda527..7bf71354f 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -45,7 +45,6 @@ from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataPars from nerfstudio.data.dataparsers.phototourism_dataparser import ( PhototourismDataParserConfig, ) -from nerfstudio.engine.dreamfusion_trainer import DreamfusionTrainerConfig from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig from nerfstudio.engine.schedulers import SchedulerConfig from nerfstudio.engine.trainer import TrainerConfig @@ -338,7 +337,7 @@ method_configs["phototourism"] = TrainerConfig( vis="viewer", ) -method_configs["dreamfusion"] = DreamfusionTrainerConfig( +method_configs["dreamfusion"] = TrainerConfig( method_name="dreamfusion", steps_per_eval_batch=50, steps_per_eval_image=50, diff --git a/nerfstudio/engine/dreamfusion_trainer.py b/nerfstudio/engine/dreamfusion_trainer.py deleted file mode 100644 index b0eb163c1..000000000 --- a/nerfstudio/engine/dreamfusion_trainer.py +++ /dev/null @@ -1,88 +0,0 @@ -# Copyright 2022 The Nerfstudio Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Dreamfusion trainer""" - -from dataclasses import dataclass, field -from typing import Dict, Tuple, Type - -import torch - -from nerfstudio.engine.trainer import Trainer, TrainerConfig -from nerfstudio.pipelines.dreamfusion_pipeline import ( - DreamfusionPipeline, - DreamfusionPipelineConfig, -) -from nerfstudio.utils import profiler - - -@dataclass -class DreamfusionTrainerConfig(TrainerConfig): - """Configuration for trainer instantiation""" - - _target: Type = field(default_factory=lambda: DreamfusionTrainer) - """target class to instantiate""" - pipeline: DreamfusionPipelineConfig = DreamfusionPipelineConfig() - """specifies the pipeline config""" - gradient_accumulation_steps: int = 1 - """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) - Trainer.__init__(self, config=config, local_rank=local_rank, world_size=world_size) - - @profiler.time_function - def train_iteration(self, step: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]: - """This function gets your training loss dict. This will be responsible for - getting the next batch of data from the DataManager and interfacing with the - Model class, feeding the data to the model's forward function. - - Args: - step: current iteration step to update sampler if using DDP (distributed) - """ - 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) - loss = loss_dict["orientation_loss"] + loss_dict["pred_normal_loss"] - if self.pipeline.config.alphas_penalty: - loss += loss_dict["alphas_loss"] * self.config.pipeline.alphas_loss_mult - if self.pipeline.config.opacity_penalty: - 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 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() - self.optimizers.scheduler_step_all(step) - - return loss, loss_dict, metrics_dict diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 58a3d2b4e..b59ce5664 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -146,7 +146,7 @@ class Trainer: 'inference': does not load any dataset into memory """ self.pipeline = self.config.pipeline.setup( - device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank + device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank, grad_scaler=self.grad_scaler ) self.optimizers = self.setup_optimizers() diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 46b479b94..3b7c45ecb 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -33,6 +33,7 @@ from rich.progress import ( TimeElapsedColumn, ) from torch import nn +from torch.cuda.amp.grad_scaler import GradScaler from torch.nn import Parameter from torch.nn.parallel import DistributedDataParallel as DDP from typing_extensions import Literal @@ -216,6 +217,7 @@ class VanillaPipeline(Pipeline): test_mode: Literal["test", "val", "inference"] = "val", world_size: int = 1, local_rank: int = 0, + grad_scaler: Optional[GradScaler] = None, ): super().__init__() self.config = config diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index 42f044a0d..13d026735 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -20,6 +20,8 @@ from typing import Optional, Type import matplotlib.pyplot as plt import numpy as np import torch +from torch.cuda.amp import custom_bwd, custom_fwd +from torch.cuda.amp.grad_scaler import GradScaler from typing_extensions import Literal from nerfstudio.data.datamanagers.dreamfusion_datamanager import ( @@ -29,6 +31,20 @@ from nerfstudio.generative.stable_diffusion import StableDiffusion from nerfstudio.pipelines.base_pipeline import VanillaPipeline, VanillaPipelineConfig +class SpecifyGradient(torch.autograd.Function): + @staticmethod + @custom_fwd + def forward(ctx, input_tensor, gt_grad): + ctx.save_for_backward(gt_grad) + return torch.zeros([], device=input_tensor.device, dtype=input_tensor.dtype) + + @staticmethod + @custom_bwd + def backward(ctx, grad): + gt_grad, = ctx.saved_tensors + batch_size = len(gt_grad) + return gt_grad / batch_size, None + @dataclass class DreamfusionPipelineConfig(VanillaPipelineConfig): """Configuration for pipeline instantiation""" @@ -78,9 +94,11 @@ class DreamfusionPipeline(VanillaPipeline): test_mode: Literal["test", "val", "inference"] = "val", world_size: int = 1, local_rank: int = 0, + grad_scaler: Optional[GradScaler] = None, ): - super().__init__(config, device, test_mode, world_size, local_rank) + super().__init__(config, device, test_mode, world_size, local_rank, grad_scaler) self.generative = True + self.grad_scaler = grad_scaler self.sd_device = ( torch.device(device) if self.config.stablediffusion_device is None @@ -162,7 +180,8 @@ class DreamfusionPipeline(VanillaPipeline): if self.config.opacity_penalty: accum_mean = torch.mean(1.0 - model_outputs["accumulation"]) sds_loss *= torch.clamp(accum_mean, min=0.5) - loss_dict["sds_loss"] = sds_loss.to(self.device) + metrics_dict["sds_loss"] = sds_loss.to(self.device) + loss_dict["latent_loss"] = SpecifyGradient.apply(self.grad_scaler.scale(latents), grad) model_outputs["latents"] = latents model_outputs["grad"] = grad diff --git a/nerfstudio/pipelines/dynamic_batch.py b/nerfstudio/pipelines/dynamic_batch.py index 293ea1706..36536eef2 100644 --- a/nerfstudio/pipelines/dynamic_batch.py +++ b/nerfstudio/pipelines/dynamic_batch.py @@ -17,9 +17,10 @@ A pipeline that dynamically chooses the number of rays to sample. """ from dataclasses import dataclass, field -from typing import Type +from typing import Optional, Type import torch +from torch.cuda.amp.grad_scaler import GradScaler from typing_extensions import Literal from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager @@ -53,6 +54,7 @@ class DynamicBatchPipeline(VanillaPipeline): test_mode: Literal["test", "val", "inference"] = "val", world_size: int = 1, local_rank: int = 0, + grad_scaler: Optional[GradScaler] = None, ): super().__init__(config, device, test_mode, world_size, local_rank) assert isinstance(