merging grad scaler pipeline addition

nerfstudio/engine/trainer.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The Nerfstudio Team. All rights reserved.
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. 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.
@@ -24,11 +24,10 @@ import time
 from dataclasses import dataclass, field
 from pathlib import Path
 from threading import Lock
-from typing import Dict, List, Literal, Optional, Tuple, Type, Union
+from typing import Dict, List, Literal, Optional, Tuple, Type, cast
 
 import torch
 from rich import box, style
-from rich.console import Console
 from rich.panel import Panel
 from rich.table import Table
 from torch.cuda.amp.grad_scaler import GradScaler
@@ -48,15 +47,13 @@ from nerfstudio.utils.decorators import (
     check_viewer_enabled,
 )
 from nerfstudio.utils.misc import step_check
+from nerfstudio.utils.rich_utils import CONSOLE
 from nerfstudio.utils.writer import EventName, TimeWriter
 from nerfstudio.viewer.server.viewer_state import ViewerState
+from nerfstudio.viewer_beta.viewer_state import ViewerState as ViewerBetaState
 
-CONSOLE = Console(width=120)
-
-TRAIN_INTERATION_OUTPUT = Tuple[  # pylint: disable=invalid-name
-    torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]
-]
-TORCH_DEVICE = Union[torch.device, str]  # pylint: disable=invalid-name
+TRAIN_INTERATION_OUTPUT = Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]
+TORCH_DEVICE = str
 
 
 @dataclass
@@ -175,15 +172,25 @@ class Trainer:
                 train_lock=self.train_lock,
             )
             banner_messages = [f"Viewer at: {self.viewer_state.viewer_url}"]
+        if self.config.is_viewer_beta_enabled() and self.local_rank == 0:
+            self.viewer_state = ViewerBetaState(
+                self.config.viewer,
+                log_filename=viewer_log_path,
+                datapath=self.base_dir,
+                pipeline=self.pipeline,
+                trainer=self,
+                train_lock=self.train_lock,
+            )
+            banner_messages = [f"Viewer Beta at: {self.viewer_state.viewer_url}"]
         self._check_viewer_warnings()
 
         self._load_checkpoint()
 
         self.callbacks = self.pipeline.get_training_callbacks(
             TrainingCallbackAttributes(
-                optimizers=self.optimizers,  # type: ignore
-                grad_scaler=self.grad_scaler,  # type: ignore
-                pipeline=self.pipeline,  # type: ignore
+                optimizers=self.optimizers,
+                grad_scaler=self.grad_scaler,
+                pipeline=self.pipeline,
             )
         )
 
@@ -223,10 +230,9 @@ class Trainer:
         """Train the model."""
         assert self.pipeline.datamanager.train_dataset is not None, "Missing DatsetInputs"
 
-        if self.pipeline.datamanager.input_data:
-            self.pipeline.datamanager.train_dataparser_outputs.save_dataparser_transform(
-                self.base_dir / "dataparser_transforms.json"
-            )
+        self.pipeline.datamanager.train_dataparser_outputs.save_dataparser_transform(
+            self.base_dir / "dataparser_transforms.json"
+        )
 
         self._init_viewer_state()
         with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME):
@@ -258,7 +264,9 @@ class Trainer:
                 if step > 1:
                     writer.put_time(
                         name=EventName.TRAIN_RAYS_PER_SEC,
-                        duration=self.pipeline.datamanager.get_train_rays_per_batch() / train_t.duration,
+                        duration=self.world_size
+                        * self.pipeline.datamanager.get_train_rays_per_batch()
+                        / train_t.duration,
                         step=step,
                         avg_over_steps=True,
                     )
@@ -304,6 +312,10 @@ class Trainer:
         table.add_row("Checkpoint Directory", str(self.checkpoint_dir))
         CONSOLE.print(Panel(table, title="[bold][green]:tada: Training Finished :tada:[/bold]", expand=False))
 
+        # after train end callbacks
+        for callback in self.callbacks:
+            callback.run_callback_at_location(step=step, location=TrainingCallbackLocation.AFTER_TRAIN)
+
         if not self.config.viewer.quit_on_train_completion:
             self._train_complete_viewer()
 
@@ -311,7 +323,7 @@ class Trainer:
     def _check_viewer_warnings(self) -> None:
         """Helper to print out any warnings regarding the way the viewer/loggers are enabled"""
         if (
-            self.config.is_viewer_enabled()
+            (self.config.is_viewer_enabled() or self.config.is_viewer_beta_enabled())
             and not self.config.is_tensorboard_enabled()
             and not self.config.is_wandb_enabled()
         ):
@@ -372,15 +384,15 @@ class Trainer:
         train_num_rays_per_batch: int = self.pipeline.datamanager.get_train_rays_per_batch()
         writer.put_time(
             name=EventName.TRAIN_RAYS_PER_SEC,
-            duration=train_num_rays_per_batch / (train_t.duration - vis_t.duration),
+            duration=self.world_size * train_num_rays_per_batch / (train_t.duration - vis_t.duration),
             step=step,
             avg_over_steps=True,
         )
 
     def _load_checkpoint(self) -> None:
         """Helper function to load pipeline and optimizer from prespecified checkpoint"""
-        load_dir: Path = self.config.load_dir
-        load_checkpoint: Path = self.config.load_checkpoint
+        load_dir = self.config.load_dir
+        load_checkpoint = self.config.load_checkpoint
         if load_dir is not None:
             load_step = self.config.load_step
             if load_step is None:
@@ -464,7 +476,7 @@ class Trainer:
                     metrics_dict[f"Gradients/{tag}"] = grad
                     total_grad += grad
 
-            metrics_dict["Gradients/Total"] = total_grad
+            metrics_dict["Gradients/Total"] = cast(torch.Tensor, total_grad)
 
         scale = self.grad_scaler.get_scale()
         self.grad_scaler.update()

nerfstudio/pipelines/base_pipeline.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The Nerfstudio Team. All rights reserved.
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. 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.
@@ -21,7 +21,7 @@ import typing
 from abc import abstractmethod
 from dataclasses import dataclass, field
 from time import time
-from typing import Any, Dict, List, Literal, Mapping, Optional, Type, Union, cast
+from typing import Any, Dict, List, Literal, Mapping, Optional, Tuple, Type, Union, cast
 
 import torch
 import torch.distributed as dist
@@ -33,14 +33,13 @@ 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 torch.cuda.amp.grad_scaler import GradScaler
 
 from nerfstudio.configs import base_config as cfg
 from nerfstudio.data.datamanagers.base_datamanager import (
     DataManager,
-    DataManagerConfig,
     VanillaDataManager,
     VanillaDataManagerConfig,
 )
@@ -91,8 +90,6 @@ class Pipeline(nn.Module):
         model: The model that will be used
     """
 
-    # pylint: disable=abstract-method
-
     datamanager: DataManager
     _model: Model
     world_size: int
@@ -187,11 +184,13 @@ class Pipeline(nn.Module):
             step: training step of the loaded checkpoint
         """
 
+    @abstractmethod
     def get_training_callbacks(
         self, training_callback_attributes: TrainingCallbackAttributes
     ) -> List[TrainingCallback]:
         """Returns the training callbacks from both the Dataloader and the Model."""
 
+    @abstractmethod
     def get_param_groups(self) -> Dict[str, List[Parameter]]:
         """Get the param groups for the pipeline.
 
@@ -206,7 +205,7 @@ class VanillaPipelineConfig(cfg.InstantiateConfig):
 
     _target: Type = field(default_factory=lambda: VanillaPipeline)
     """target class to instantiate"""
-    datamanager: DataManagerConfig = VanillaDataManagerConfig()
+    datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig()
     """specifies the datamanager config"""
     model: ModelConfig = ModelConfig()
     """specifies the model config"""
@@ -224,6 +223,7 @@ class VanillaPipeline(Pipeline):
             'inference': does not load any dataset into memory
         world_size: total number of machines available
         local_rank: rank of current machine
+        grad_scaler: gradient scaler used in the trainer
 
     Attributes:
         datamanager: The data manager that will be used
@@ -278,7 +278,7 @@ class VanillaPipeline(Pipeline):
             step: current iteration step to update sampler if using DDP (distributed)
         """
         ray_bundle, batch = self.datamanager.next_train(step)
-        model_outputs = self.model(ray_bundle)
+        model_outputs = self._model(ray_bundle)  # train distributed data parallel model if world_size > 1
         metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
 
         if self.config.datamanager.camera_optimizer is not None:
@@ -304,7 +304,7 @@ class VanillaPipeline(Pipeline):
         raise NotImplementedError
 
     @profiler.time_function
-    def get_eval_loss_dict(self, step: int):
+    def get_eval_loss_dict(self, step: int) -> Tuple[Any, Dict[str, Any], Dict[str, Any]]:
         """This function gets your evaluation loss dict. It needs to get the data
         from the DataManager and feed it to the model's forward function
 

nerfstudio/pipelines/dynamic_batch.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The Nerfstudio Team. All rights reserved.
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. 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.
@@ -40,8 +40,6 @@ class DynamicBatchPipelineConfig(VanillaPipelineConfig):
 class DynamicBatchPipeline(VanillaPipeline):
     """Pipeline with logic for changing the number of rays per batch."""
 
-    # pylint: disable=abstract-method
-
     config: DynamicBatchPipelineConfig
     datamanager: VanillaDataManager
     dynamic_num_rays_per_batch: int
@@ -86,11 +84,12 @@ class DynamicBatchPipeline(VanillaPipeline):
                 "'num_samples_per_batch' is not in metrics_dict."
                 "Please return 'num_samples_per_batch' in the models get_metrics_dict function to use this method."
             )
-        self._update_dynamic_num_rays_per_batch(metrics_dict["num_samples_per_batch"])
+        self._update_dynamic_num_rays_per_batch(int(metrics_dict["num_samples_per_batch"]))
         self._update_pixel_samplers()
 
         # add the number of rays
         assert "num_rays_per_batch" not in metrics_dict
+        assert self.datamanager.train_pixel_sampler is not None
         metrics_dict["num_rays_per_batch"] = torch.tensor(self.datamanager.train_pixel_sampler.num_rays_per_batch)
 
         return model_outputs, loss_dict, metrics_dict
@@ -100,6 +99,7 @@ class DynamicBatchPipeline(VanillaPipeline):
 
         # add the number of rays
         assert "num_rays_per_batch" not in metrics_dict
+        assert self.datamanager.eval_pixel_sampler is not None
         metrics_dict["num_rays_per_batch"] = torch.tensor(self.datamanager.eval_pixel_sampler.num_rays_per_batch)
 
         return model_outputs, loss_dict, metrics_dict