diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index efc9e170c..e786b9ab2 100644 --- a/nerfstudio/engine/trainer.py +++ b/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() diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 40f707822..1e1e0fab5 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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 diff --git a/nerfstudio/pipelines/dynamic_batch.py b/nerfstudio/pipelines/dynamic_batch.py index 1202e6503..65679d001 100644 --- a/nerfstudio/pipelines/dynamic_batch.py +++ b/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