diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 03e20fada..5e0263ef8 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -34,7 +34,7 @@ from nerfactory.models.base import Model from nerfactory.models.nerfw import NerfWModel from nerfactory.models.tensorf import TensoRFModel from nerfactory.optimizers.schedulers import ExponentialDecaySchedule -from nerfactory.pipelines.base import Pipeline +from nerfactory.pipelines.base import VanillaPipeline from nerfactory.utils import writer @@ -296,10 +296,10 @@ class TensoRFModelConfig(VanillaModelConfig): # Pipeline related configs @dataclass -class PipelineConfig(InstantiateConfig): +class VanillaPipelineConfig(InstantiateConfig): """Configuration for pipeline instantiation""" - _target: Type = Pipeline + _target: Type = VanillaPipeline """target class to instantiate""" datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig() """specifies the datamanager config""" @@ -390,7 +390,7 @@ class Config(PrintableConfig): logging: LoggingConfig = LoggingConfig() viewer: ViewerConfig = ViewerConfig() trainer: TrainerConfig = TrainerConfig() - pipeline: PipelineConfig = PipelineConfig() + pipeline: VanillaPipelineConfig = VanillaPipelineConfig() optimizers: Dict[str, Any] = to_immutable_dict( { "fields": { diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py index cf5934c8a..6bc0dff48 100644 --- a/nerfactory/configs/base_configs.py +++ b/nerfactory/configs/base_configs.py @@ -29,13 +29,13 @@ from nerfactory.configs.base import ( LoggingConfig, MipNerf360DataParserConfig, NerfWModelConfig, - PipelineConfig, RAdamOptimizerConfig, SchedulerConfig, TensoRFModelConfig, TrainerConfig, VanillaDataManagerConfig, VanillaModelConfig, + VanillaPipelineConfig, ViewerConfig, ) from nerfactory.datamanagers.dataparsers.friends_parser import FriendsDataParserConfig @@ -50,7 +50,7 @@ base_configs: Dict[str, Config] = {} base_configs["compound"] = Config( method_name="compound", trainer=TrainerConfig(mixed_precision=True), - pipeline=PipelineConfig( + pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig(dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192), model=CompoundModelConfig(eval_num_rays_per_chunk=8192), ), @@ -67,7 +67,7 @@ base_configs["compound"] = Config( base_configs["instant-ngp"] = Config( method_name="instant-ngp", trainer=TrainerConfig(steps_per_eval_batch=500, steps_per_save=2000, mixed_precision=True), - pipeline=PipelineConfig( + pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig(dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192), model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192), ), @@ -84,7 +84,7 @@ base_configs["instant-ngp"] = Config( base_configs["mipnerf-360"] = Config( method_name="mipnerf-360", trainer=TrainerConfig(steps_per_eval_batch=200), - pipeline=PipelineConfig( + pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig(dataparser=MipNerf360DataParserConfig(), train_num_rays_per_batch=8192), model=VanillaModelConfig( _target=MipNerf360Model, @@ -105,7 +105,7 @@ base_configs["mipnerf-360"] = Config( base_configs["mipnerf"] = Config( method_name="mipnerf", - pipeline=PipelineConfig( + pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig(dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192), model=VanillaModelConfig( _target=MipNerfModel, @@ -125,7 +125,7 @@ base_configs["mipnerf"] = Config( base_configs["nerfw"] = Config( method_name="nerfw", - pipeline=PipelineConfig( + pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=FriendsDataParserConfig(), ), @@ -142,7 +142,7 @@ base_configs["nerfw"] = Config( base_configs["semantic-nerf"] = Config( method_name="semantic-nerf", - pipeline=PipelineConfig( + pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=FriendsDataParserConfig(), ), @@ -163,7 +163,7 @@ base_configs["semantic-nerf"] = Config( base_configs["vanilla-nerf"] = Config( method_name="vanilla-nerf", - pipeline=PipelineConfig( + pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=BlenderDataParserConfig(), ), @@ -180,7 +180,7 @@ base_configs["vanilla-nerf"] = Config( base_configs["tensorf"] = Config( method_name="tensorf", trainer=TrainerConfig(mixed_precision=True), - pipeline=PipelineConfig( + pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=BlenderDataParserConfig(), ), diff --git a/nerfactory/datamanagers/base.py b/nerfactory/datamanagers/base.py index a2463abd9..c67aef23e 100644 --- a/nerfactory/datamanagers/base.py +++ b/nerfactory/datamanagers/base.py @@ -28,7 +28,7 @@ from torch.utils.data.distributed import DistributedSampler from nerfactory.cameras.rays import RayBundle from nerfactory.configs import base as cfg from nerfactory.datamanagers.dataloaders import ( - CacheImageDataloader, + CacheDataloader, FixedIndicesEvalDataloader, RandIndicesEvalDataloader, ) @@ -246,7 +246,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method ) else: sampler = None - self.train_image_dataloader = CacheImageDataloader( + self.train_image_dataloader = CacheDataloader( self.train_input_dataset, num_images_to_sample_from=self.config.train_num_images_to_sample_from, device=self.device, @@ -267,7 +267,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method ) else: sampler = None - self.eval_image_dataloader = CacheImageDataloader( + self.eval_image_dataloader = CacheDataloader( self.eval_input_dataset, num_images_to_sample_from=self.config.eval_num_images_to_sample_from, device=self.device, diff --git a/nerfactory/datamanagers/dataloaders.py b/nerfactory/datamanagers/dataloaders.py index 2520beef3..fb7b0dd02 100644 --- a/nerfactory/datamanagers/dataloaders.py +++ b/nerfactory/datamanagers/dataloaders.py @@ -21,7 +21,7 @@ from abc import abstractmethod from typing import Dict, Optional, Tuple, Union import torch -from torch.utils.data import default_collate +from torch.utils.data import Dataset, default_collate from torch.utils.data.dataloader import DataLoader from nerfactory.cameras.cameras import Cameras @@ -30,8 +30,8 @@ from nerfactory.datamanagers.datasets import InputDataset from nerfactory.utils.misc import get_dict_to_torch -class CacheImageDataloader(DataLoader): - """Collated image dataset that implements caching of images. +class CacheDataloader(DataLoader): + """Collated image dataset that implements caching of default-pytorch-collatable data. Creates batches of the InputDataset return type. Args: @@ -43,7 +43,7 @@ class CacheImageDataloader(DataLoader): def __init__( self, - dataset: InputDataset, + dataset: Dataset, num_images_to_sample_from: int = -1, num_times_to_repeat_images: int = 0, device: Union[torch.device, str] = "cpu", diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 457b39f46..a4f9329d9 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -31,7 +31,7 @@ from torch.cuda.amp.grad_scaler import GradScaler from nerfactory.configs import base as cfg from nerfactory.optimizers.optimizers import Optimizers, setup_optimizers -from nerfactory.pipelines.base import Pipeline +from nerfactory.pipelines.base import VanillaPipeline from nerfactory.utils import profiler, writer from nerfactory.utils.callbacks import ( TrainingCallback, @@ -79,7 +79,7 @@ class Trainer: callbacks: The callbacks object. """ - pipeline: Pipeline + pipeline: VanillaPipeline optimizers: Optimizers callbacks: List[TrainingCallback] @@ -119,7 +119,7 @@ class Trainer: Args: test_mode: Whether to setup for testing. Defaults to False. """ - self.pipeline: Pipeline = self.config.pipeline.setup( + self.pipeline: VanillaPipeline = self.config.pipeline.setup( device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank ) self.optimizers = setup_optimizers(self.config.optimizers, self.pipeline.get_param_groups()) diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index cd7ae2638..3d71f57b7 100644 --- a/nerfactory/pipelines/base.py +++ b/nerfactory/pipelines/base.py @@ -36,7 +36,7 @@ from torch.nn import Parameter from torch.nn.parallel import DistributedDataParallel as DDP from nerfactory.configs import base as cfg -from nerfactory.datamanagers.base import DataManager +from nerfactory.datamanagers.base import DataManager, VanillaDataManager from nerfactory.models.base import Model from nerfactory.utils import profiler from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes @@ -69,19 +69,6 @@ class Pipeline(nn.Module): simplifying the model classes, which each may have different forward() methods and so on. - - TODO: For viewer functionality to be added down the line, we should make sure - that we are still abstracting away the Model from the end user. The viewer function - should probably be done by adding a new iterator on the base data manager that will - talk to the actual viewer. This is probably ideal to have the viewer be - primarily located in the data manager (first because it makes sense as the - viewers main job in this context is to feed in data for the model to load) - so that we can have an easier time ensuring that the viewer is always - returning the same formatted data as for in train / eval. All this is pending changes to - be done in the future... but just bear in mind that if learned parameters are in the data manager, - the viewer may have to use those parameters as well. - - Args: config: configuration to instantiate pipeline device: location to place model and data @@ -94,11 +81,115 @@ class Pipeline(nn.Module): self.model: The model that will be used """ + datamanager: DataManager + model: Model + + @property + def device(self): + """Returns the device that the model is on.""" + return module_wrapper(self.model).device + + @profiler.time_function + def get_train_loss_dict(self, step: int): + """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) + """ + if self.world_size > 1 and step: + assert self.datamanager.train_sampler is not None + self.datamanager.train_sampler.set_epoch(step) + ray_bundle, batch = self.datamanager.next_train(step) + model_outputs = self.model(ray_bundle, batch) + metrics_dict = self.model.get_metrics_dict(model_outputs, batch) + loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict) + + return model_outputs, loss_dict, metrics_dict + + @profiler.time_function + def get_eval_loss_dict(self, step: int): + """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 + + Args: + step: current iteration step + """ + self.eval() + if self.world_size > 1: + assert self.datamanager.eval_sampler is not None + self.datamanager.eval_sampler.set_epoch(step) + ray_bundle, batch = self.datamanager.next_eval(step) + model_outputs = self.model(ray_bundle, batch) + loss_dict = self.model.get_loss_dict(model_outputs, batch) + metrics_dict = self.model.get_metrics_dict(model_outputs, batch) + self.train() + return model_outputs, loss_dict, metrics_dict + + @abstractmethod + @profiler.time_function + def get_eval_image_metrics_and_images(self, step: int): + """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 + + Args: + step: current iteration step + """ + + @abstractmethod + @profiler.time_function + def get_average_eval_image_metrics(self, step: Optional[int] = None): + """Iterate over all the images in the eval dataset and get the average.""" + + @abstractmethod + def log_test_image_outputs(self) -> None: + """Log the test image outputs""" + + def load_pipeline(self, loaded_state: Dict[str, Any]) -> None: + """Load the checkpoint from the given path + + Args: + loaded_state: pre-trained model state dict + """ + + def get_training_callbacks( + self, training_callback_attributes: TrainingCallbackAttributes + ) -> List[TrainingCallback]: + """Returns the training callbacks from both the Dataloader and the Model.""" + + def get_param_groups(self) -> Dict[str, List[Parameter]]: + """Get the param groups for the pipeline. + + Returns: + A list of dictionaries containing the pipeline's param groups. + """ + + +class VanillaPipeline(Pipeline): + """The pipeline class for the vanilla nerf setup of multiple cameras for one or a few scenes. + + config: configuration to instantiate pipeline + device: location to place model and data + test_mode: if True, loads test datset. if False, loads train/eval datasets + world_size: total number of machines available + local_rank: rank of current machine + + Attributes: + self.datamanager: The data manager that will be used + self.model: The model that will be used + """ + def __init__( - self, config: cfg.PipelineConfig, device: str, test_mode: bool = False, world_size: int = 1, local_rank: int = 0 + self, + config: cfg.VanillaPipelineConfig, + device: str, + test_mode: bool = False, + world_size: int = 1, + local_rank: int = 0, ): super().__init__() - self.datamanager: DataManager = config.datamanager.setup( + self.datamanager: VanillaDataManager = config.datamanager.setup( device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank ) self.datamanager.to(device)