Remove Dreamfusion Trainer (#1423)

* Dreamfusion trainer no longer needed

* Added typing for grad scaler

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,

nerfstudio/engine/dreamfusion_trainer.py

@@ -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

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()
 

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

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

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(