tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

feat(engine): improve type hints (#1509)

Browse Source 
* 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
This commit is contained in:
Saurav Maheshkar 2023-02-26 00:33:43 +00:00 committed by GitHub
parent 5158cc36ff
commit 0b87c43d0d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 33 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
nerfstudio/engine/callbacks.py
View File

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

6
nerfstudio/engine/optimizers.py
View File

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

2
nerfstudio/engine/schedulers.py
View File

@ -40,7 +40,7 @@ class Scheduler:
config: SchedulerConfig
def __init__(self, config: SchedulerConfig):
def __init__(self, config: SchedulerConfig) -> None:
super().__init__()
self.config = config

49
nerfstudio/engine/trainer.py
View File

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