From 0b87c43d0d5cf4b41cd4e9541b7f9f0f8efe7e08 Mon Sep 17 00:00:00 2001 From: Saurav Maheshkar Date: Sun, 26 Feb 2023 00:33:43 +0000 Subject: [PATCH] feat(engine): improve type hints (#1509) * feat(engine): improve type hints * fix: add pylint ignore statement * fix: order imports as per isort * fix: update Optimizer setup return type * feat: add new TORCH_DEVICE type --- nerfstudio/engine/callbacks.py | 4 +-- nerfstudio/engine/optimizers.py | 6 ++-- nerfstudio/engine/schedulers.py | 2 +- nerfstudio/engine/trainer.py | 49 ++++++++++++++++++--------------- 4 files changed, 33 insertions(+), 28 deletions(-) diff --git a/nerfstudio/engine/callbacks.py b/nerfstudio/engine/callbacks.py index 3c7219d67..6be1d8097 100644 --- a/nerfstudio/engine/callbacks.py +++ b/nerfstudio/engine/callbacks.py @@ -79,7 +79,7 @@ class TrainingCallback: self.args = args if args is not None else [] self.kwargs = kwargs if kwargs is not None else {} - def run_callback(self, step: int): + def run_callback(self, step: int) -> None: """Callback to run after training step Args: @@ -92,7 +92,7 @@ class TrainingCallback: if step in self.iters: self.func(*self.args, **self.kwargs, step=step) - def run_callback_at_location(self, step: int, location: TrainingCallbackLocation): + def run_callback_at_location(self, step: int, location: TrainingCallbackLocation) -> None: """Runs the callback if it's supposed to be run at the given location. Args: diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py index 786c51115..d2cdb022a 100644 --- a/nerfstudio/engine/optimizers.py +++ b/nerfstudio/engine/optimizers.py @@ -44,7 +44,7 @@ class OptimizerConfig(base_config.PrintableConfig): # TODO: somehow make this more generic. i dont like the idea of overriding the setup function # but also not sure how to go about passing things into predefined torch objects. - def setup(self, params) -> Any: + def setup(self, params) -> torch.optim.Optimizer: """Returns the instantiated object using the config.""" kwargs = vars(self).copy() kwargs.pop("_target") @@ -76,7 +76,7 @@ class Optimizers: param_groups: A dictionary of parameter groups to optimize. """ - def __init__(self, config: Dict[str, Any], param_groups: Dict[str, List[Parameter]]): + def __init__(self, config: Dict[str, Any], param_groups: Dict[str, List[Parameter]]) -> None: self.config = config self.optimizers = {} self.schedulers = {} @@ -127,7 +127,7 @@ class Optimizers: torch.nn.utils.clip_grad_norm_(self.parameters[param_group], max_norm) grad_scaler.step(optimizer) - def optimizer_step_all(self): + def optimizer_step_all(self) -> None: """Run step for all optimizers.""" for param_group, optimizer in self.optimizers.items(): # note that they key is the parameter name diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py index d4554c7d8..0b6586ffc 100644 --- a/nerfstudio/engine/schedulers.py +++ b/nerfstudio/engine/schedulers.py @@ -40,7 +40,7 @@ class Scheduler: config: SchedulerConfig - def __init__(self, config: SchedulerConfig): + def __init__(self, config: SchedulerConfig) -> None: super().__init__() self.config = config diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index b3e786bd3..f900dd863 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -23,7 +23,7 @@ import os import time from dataclasses import dataclass, field from pathlib import Path -from typing import Dict, List, Optional, Tuple, Type +from typing import Dict, List, Optional, Tuple, Type, Union import torch from rich.console import Console @@ -50,6 +50,11 @@ from nerfstudio.viewer.server import viewer_utils 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 + @dataclass class TrainerConfig(ExperimentConfig): @@ -104,25 +109,25 @@ class Trainer: optimizers: Optimizers callbacks: List[TrainingCallback] - def __init__(self, config: TrainerConfig, local_rank: int = 0, world_size: int = 1): + def __init__(self, config: TrainerConfig, local_rank: int = 0, world_size: int = 1) -> None: self.config = config self.local_rank = local_rank self.world_size = world_size - self.device = "cpu" if world_size == 0 else f"cuda:{local_rank}" - self.mixed_precision = self.config.mixed_precision + self.device: TORCH_DEVICE = "cpu" if world_size == 0 else f"cuda:{local_rank}" + self.mixed_precision: bool = self.config.mixed_precision if self.device == "cpu": self.mixed_precision = False CONSOLE.print("Mixed precision is disabled for CPU training.") - self._start_step = 0 + self._start_step: int = 0 # optimizers self.grad_scaler = GradScaler(enabled=self.mixed_precision) - self.base_dir = config.get_base_dir() + self.base_dir: Path = config.get_base_dir() # directory to save checkpoints - self.checkpoint_dir = config.get_checkpoint_dir() + self.checkpoint_dir: Path = config.get_checkpoint_dir() CONSOLE.log(f"Saving checkpoints to: {self.checkpoint_dir}") # set up viewer if enabled - viewer_log_path = self.base_dir / config.viewer.relative_log_filename + viewer_log_path: Path = self.base_dir / config.viewer.relative_log_filename self.viewer_state, banner_messages = None, None if self.config.is_viewer_enabled() and local_rank == 0: self.viewer_state, banner_messages = viewer_utils.setup_viewer( @@ -130,13 +135,13 @@ class Trainer: ) self._check_viewer_warnings() # set up writers/profilers if enabled - writer_log_path = self.base_dir / config.logging.relative_log_dir + writer_log_path: Path = self.base_dir / config.logging.relative_log_dir writer.setup_event_writer(config.is_wandb_enabled(), config.is_tensorboard_enabled(), log_dir=writer_log_path) writer.setup_local_writer(config.logging, max_iter=config.max_num_iterations, banner_messages=banner_messages) writer.put_config(name="config", config_dict=dataclasses.asdict(config), step=0) profiler.setup_profiler(config.logging) - def setup(self, test_mode: Literal["test", "val", "inference"] = "val"): + def setup(self, test_mode: Literal["test", "val", "inference"] = "val") -> None: """Setup the Trainer by calling other setup functions. Args: @@ -243,7 +248,7 @@ class Trainer: self._always_render(step) @check_main_thread - def _always_render(self, step): + def _always_render(self, step: int) -> None: if self.config.is_viewer_enabled(): while True: self.viewer_state.vis["renderingState/isTraining"].write(False) @@ -257,7 +262,7 @@ class Trainer: and not self.config.is_tensorboard_enabled() and not self.config.is_wandb_enabled() ): - string = ( + string: str = ( "[NOTE] Not running eval iterations since only viewer is enabled.\n" "Use [yellow]--vis {wandb, tensorboard, viewer+wandb, viewer+tensorboard}[/yellow] to run with eval." ) @@ -275,7 +280,7 @@ class Trainer: self._always_render(self._start_step) @check_viewer_enabled - def _update_viewer_state(self, step: int): + def _update_viewer_state(self, step: int) -> None: """Updates the viewer state by rendering out scene with current pipeline Returns the time taken to render scene. @@ -284,7 +289,7 @@ class Trainer: """ assert self.viewer_state is not None with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _: - num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch + num_rays_per_batch: int = self.config.pipeline.datamanager.train_num_rays_per_batch try: self.viewer_state.update_scene(self, step, self.pipeline.model, num_rays_per_batch) except RuntimeError: @@ -295,7 +300,7 @@ class Trainer: ) @check_viewer_enabled - def _update_viewer_rays_per_sec(self, train_t: TimeWriter, vis_t: TimeWriter, step: int): + def _update_viewer_rays_per_sec(self, train_t: TimeWriter, vis_t: TimeWriter, step: int) -> None: """Performs update on rays/sec calculation for training Args: @@ -303,7 +308,7 @@ class Trainer: vis_t: timer object carrying time to execute visualization step step: current step """ - train_num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch + train_num_rays_per_batch: int = self.config.pipeline.datamanager.train_num_rays_per_batch writer.put_time( name=EventName.TRAIN_RAYS_PER_SEC, duration=train_num_rays_per_batch / (train_t.duration - vis_t.duration), @@ -313,14 +318,14 @@ class Trainer: def _load_checkpoint(self) -> None: """Helper function to load pipeline and optimizer from prespecified checkpoint""" - load_dir = self.config.load_dir + load_dir: Path = self.config.load_dir if load_dir is not None: load_step = self.config.load_step if load_step is None: print("Loading latest checkpoint from load_dir") # NOTE: this is specific to the checkpoint name format load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(load_dir))[-1] - load_path = load_dir / f"step-{load_step:09d}.ckpt" + load_path: Path = load_dir / f"step-{load_step:09d}.ckpt" assert load_path.exists(), f"Checkpoint {load_path} does not exist" loaded_state = torch.load(load_path, map_location="cpu") self._start_step = loaded_state["step"] + 1 @@ -343,7 +348,7 @@ class Trainer: if not self.checkpoint_dir.exists(): self.checkpoint_dir.mkdir(parents=True, exist_ok=True) # save the checkpoint - ckpt_path = self.checkpoint_dir / f"step-{step:09d}.ckpt" + ckpt_path: Path = self.checkpoint_dir / f"step-{step:09d}.ckpt" torch.save( { "step": step, @@ -363,14 +368,14 @@ class Trainer: f.unlink() @profiler.time_function - def train_iteration(self, step: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]: + def train_iteration(self, step: int) -> TRAIN_INTERATION_OUTPUT: """Run one iteration with a batch of inputs. Returns dictionary of model losses. Args: step: Current training step. """ self.optimizers.zero_grad_all() - cpu_or_cuda_str = self.device.split(":")[0] + cpu_or_cuda_str: str = self.device.split(":")[0] with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision): _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step) loss = functools.reduce(torch.add, loss_dict.values()) @@ -396,7 +401,7 @@ class Trainer: @check_eval_enabled @profiler.time_function - def eval_iteration(self, step): + def eval_iteration(self, step: int) -> None: """Run one iteration with different batch/image/all image evaluations depending on step size. Args: