diff --git a/docs/reference/api/data.rst b/docs/reference/api/data.rst index f08638e34..b3759b89c 100644 --- a/docs/reference/api/data.rst +++ b/docs/reference/api/data.rst @@ -3,6 +3,13 @@ Data ============ +Data Parsers +---------------- + +.. automodule:: nerfactory.dataloaders.data_parsers + :members: + :show-inheritance: + Image Dataset ---------------- @@ -10,10 +17,10 @@ Image Dataset :members: :show-inheritance: -Image Sampler +Image Dataloader ---------------- -.. automodule:: nerfactory.dataloaders.image_sampler +.. automodule:: nerfactory.dataloaders.image_dataloader :members: :show-inheritance: diff --git a/docs/tutorials/pipelines/creating_pipelines.ipynb b/docs/tutorials/pipelines/creating_pipelines.ipynb index 9ab6c250a..b9fad4c24 100644 --- a/docs/tutorials/pipelines/creating_pipelines.ipynb +++ b/docs/tutorials/pipelines/creating_pipelines.ipynb @@ -10,15 +10,15 @@ "\n", "> Our goal is to enable Nerfactory users to implement an entire NeRF method by only editing one file.\n", "\n", - "Here we explain how to create custom pipelines with nerfactory. Pipelines are composed of two components, namely a Dataloader and a Model. The pipeline is responsible for routing RayBundle instances returned from the Dataloader into the Model.\n", + "Here we explain how to create custom pipelines with nerfactory. Pipelines are composed of two components, namely a DataManager and a Model. The pipeline is responsible for routing RayBundle instances returned from the DataManager into the Model.\n", "\n", - "A NeRF method is implemented as a Pipeline. Users can create a Pipeline with a Dataloader and Model. We provide default implementations for these components prefixed with the name Vanilla (e.g., `VanillaPipeline`, `VanillaDataloader` which inherit from the `Pipeline` and `Dataloader` class, respectively) and specific Models (e.g., `InstantNGPModel` or `TensoRFModel` which inherit from the `Model` class). With most NeRF methods, `VanillaPipeline` and `VanillaDataloader` can be used and only the model needs to be changed. We provide many already implmented Models for people to get started with 3D reconstruction. However, in this tutorial we will demonstrate how to customize each of the Pipeline, Dataloader, and Model with an example.\n", + "A NeRF method is implemented as a Pipeline. Users can create a Pipeline with a DataManager and Model. We provide default implementations for these components prefixed with the name Vanilla (e.g., `VanillaPipeline`, `VanillaDataManager` which inherit from the `Pipeline` and `DataManager` class, respectively) and specific Models (e.g., `InstantNGPModel` or `TensoRFModel` which inherit from the `Model` class). With most NeRF methods, `VanillaPipeline` and `VanillaDataManager` can be used and only the model needs to be changed. We provide many already implmented Models for people to get started with 3D reconstruction. However, in this tutorial we will demonstrate how to customize each of the Pipeline, DataManager, and Model with an example.\n", "\n", - "We'll show how to make an incremental dataloader where we progressively add cameras to train with from multiple scenes and importance sample rays from pixels with a high loss. The extensible features that we show in this tutorial are the following:\n", - "- a Dataloader that incrementally adds cameras\n", - "- a Dataloader that uses multiple scenes\n", + "We'll show how to make an incremental DataManager where we progressively add cameras to train with from multiple scenes and importance sample rays from pixels with a high loss. The extensible features that we show in this tutorial are the following:\n", + "- a DataManager that incrementally adds cameras\n", + "- a DataManager that uses multiple scenes\n", "- a Model that is conditioned on the scene\n", - "- a Pipeline that modifies the Dataloader's PixelSampler to sample more rays in regions with high loss" + "- a Pipeline that modifies the DataManager's PixelSampler to sample more rays in regions with high loss" ] }, { @@ -62,34 +62,29 @@ "
  • Spatial Distortions
    • \n", "\n", "\n", - "#### Dataloader Components\n", + "#### DataManager Components\n", "\n", - "> The Dataloader provides a RayBundle (3D line segments with metadata needed to run the forward pass) and a dictionary of information needed at train-time to supervise the pipeline components.\n", + "> The DataManager provides a RayBundle (3D line segments with metadata needed to run the forward pass) and a dictionary of information needed at train-time to supervise the pipeline components.\n", "\n", - "Similar to the Model components, we provide Dataloader components that are strung together to generate RayBundle instances. A RayBundle contains everything needed to generate a ray (e.g., origin and direction, near and far planes) as well as any additional metadata that the Model requires (e.g., camera index, scene index, etc.).\n", + "Similar to the Model components, we provide DataManager components that are strung together to generate RayBundle instances. A RayBundle contains everything needed to generate a ray (e.g., origin and direction, near and far planes) as well as any additional metadata that the Model requires (e.g., camera index, scene index, etc.).\n", "\n", "
    \n", - "Dataloader Components\n", + "DataManager Components\n", "
    \n", - "Scene Datasets:
    \n", + "Datasets Parsers:
    \n", " Takes a data location and returns a DatasetInputs instance. A DatasetInputs will contain all data needed throughout the pipeline.\n", - "
  • Vanilla Image Dataset
    • \n", - "
  • Panoptic Image Dataset
    • \n", + "
  • Vanilla Image DataParser
    • \n", + "
  • Panoptic Image DataParser
    • \n", "
    \n", "Image Datasets:
    \n", " Takes image filename information and can return images.\n", "
  • Vanilla Image Dataset
    • \n", "
  • Panoptic Image Dataset
    • \n", "
    \n", - "Image Samplers:
    \n", + "Image Dataloader:
    \n", " Takes an Image Dataset and samples images from it.\n", - "
  • Vanilla Image Sampler
    • \n", - "
  • Cache Image Sampler
    • \n", - "
    \n", - "Pixel Samplers:
    \n", - " Takes sampled images and returns pixel locations.\n", - "
  • Vanilla Image Sampler
    • \n", - "
  • Cache Image Sampler
    • \n", + "
  • Vanilla Image Dataloader
    • \n", + "
  • Cache Image Dataloader
    • \n", "
    \n", "Pixel Samplers:
    \n", " Takes sampled images and returns pixel locations.\n", @@ -123,8 +118,8 @@ "from typing import Dict, Optional, Tuple, List\n", "import torch\n", "\n", - "from nerfactory.utils.config import DataloaderConfig, ModelConfig, PipelineConfig, ViewerConfig\n", - "from nerfactory.dataloaders.base import Dataloader\n", + "from nerfactory.utils.config import DataManagerConfig, ModelConfig, PipelineConfig, ViewerConfig\n", + "from nerfactory.dataloaders.base import DataManager\n", "from nerfactory.dataloaders.structs import DatasetInputs\n", "from nerfactory.models.base import Model\n", "from nerfactory.pipelines.base import Pipeline\n", @@ -140,17 +135,17 @@ }, "outputs": [], "source": [ - "class CustomDataloaderConfig(DataloaderConfig):\n", - " _target = \"CustomDataloader\"\n", + "class CustomDataManagerConfig(DataManagerConfig):\n", + " _target = \"CustomDataManager\"\n", " scene_names: List[str] = [\"chair\", \"drums\", \"fern\", \"lego\"]\n", " steps_per_add_cameras: int = 100 # every X iterations, add more cameras\n", " num_cameras_per_add: int = 10 # number of cameras to add\n", "\n", "\n", - "class CustomDataloader(Dataloader):\n", + "class CustomDataManager(DataManager):\n", " \"\"\"A custom dataloader that incrementally adds cameras and has multiple scenes.\"\"\"\n", "\n", - " config: CustomDataloaderConfig\n", + " config: CustomDataManagerConfig\n", "\n", " def populate_train_modules(self):\n", " \"\"\"Populate the train dataloader modules.\"\"\"\n", @@ -159,7 +154,7 @@ " for scene_name in self.config.scene_names:\n", " dataset_inputs: DatasetInputs = None # TODO\n", " image_dataset = ImageDataset(dataset_inputs)\n", - " image_sampler = CacheImageSampler(image_dataset=image_dataset) # a torch dataloader\n", + " image_sampler = CacheImageDataloader(image_dataset=image_dataset) # a torch dataloader\n", " self.scene_name_to_image_sampler[scene_name] = image_sampler\n", " self.train_pixel_sampler = PixelSampler\n", " self.train_ray_generator = RayGenerator # nn.Module\n", @@ -250,7 +245,7 @@ }, "outputs": [], "source": [ - "dataloader = Dataloader()\n", + "dataloader = DataManager()\n", "model = SceneConditionNGPModel()\n", "pipeline = CustomPipeline.from_dataloader_and_model(dataloader=dataloader, model=model)" ] @@ -352,7 +347,7 @@ }, "vscode": { "interpreter": { - "hash": "e691f033e0f2f1b9c0a11e7e81375a1e27aec87a71fd1b1eaa545f7a5e61b3f6" + "hash": "100252972d4ee5286ab455ef1c778abde05b4affca9d1be78633d3797abf7e8c" } } }, diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 5012c4aaf..3986547ea 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -24,10 +24,10 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Type import torch from nerfactory.configs.utils import to_immutable_dict -from nerfactory.dataloaders.base import VanillaDataloader -from nerfactory.dataloaders.datasets import ( +from nerfactory.dataloaders.base import VanillaDataManager +from nerfactory.dataloaders.data_parsers import ( Blender, - Dataset, + DataParser, Friends, InstantNGP, Mipnerf360, @@ -162,14 +162,14 @@ class TrainerConfig(PrintableConfig): # Dataset related configs @dataclass -class DatasetConfig(InstantiateConfig): +class DataParserConfig(InstantiateConfig): """Basic dataset config""" - _target: Type = Dataset + _target: Type = DataParser @dataclass -class BlenderDatasetConfig(DatasetConfig): +class BlenderDataParserConfig(DataParserConfig): """Blender dataset config""" _target: Type = Blender @@ -180,7 +180,7 @@ class BlenderDatasetConfig(DatasetConfig): @dataclass -class FriendsDatasetConfig(DatasetConfig): +class FriendsDataParserConfig(DataParserConfig): """Friends dataset config""" _target: Type = Friends @@ -188,7 +188,7 @@ class FriendsDatasetConfig(DatasetConfig): @dataclass -class MipNerf360DatasetConfig(DatasetConfig): +class MipNerf360DataParserConfig(DataParserConfig): """Mipnerf 360 dataset config""" _target: Type = Mipnerf360 @@ -200,7 +200,18 @@ class MipNerf360DatasetConfig(DatasetConfig): @dataclass -class Record3DDatasetConfig(DatasetConfig): +class InstantNGPDataParserConfig(DataParserConfig): + """Mipnerf 360 dataset config""" + + _target: Type = InstantNGP + data_directory: Path = Path("data/ours/posterv2") + scale_factor: float = 1.0 + downscale_factor: int = 1 + scene_scale: float = 0.33 + + +@dataclass +class Record3DDataParserConfig(DataParserConfig): """Mipnerf 360 dataset config""" _target: Type = Mipnerf360 @@ -212,11 +223,11 @@ class Record3DDatasetConfig(DatasetConfig): @dataclass -class VanillaDataloaderConfig(InstantiateConfig): - """Configuration for dataloader instantiation""" +class VanillaDataManagerConfig(InstantiateConfig): + """Configuration for data manager instantiation""" - _target: Type = VanillaDataloader - train_dataset: DatasetConfig = BlenderDatasetConfig() + _target: Type = VanillaDataManager + train_dataset: DataParserConfig = BlenderDataParserConfig() image_dataset_type: str = "rgb" train_num_rays_per_batch: int = 1024 train_num_images_to_sample_from: int = -1 @@ -226,25 +237,14 @@ class VanillaDataloaderConfig(InstantiateConfig): @dataclass -class FriendsDataloaderConfig(VanillaDataloaderConfig): - """Friends dataloader config""" +class FriendsDataManagerConfig(VanillaDataManagerConfig): + """Friends data manager config""" - _target: Type = VanillaDataloader - train_dataset: DatasetConfig = FriendsDatasetConfig() + _target: Type = VanillaDataManager + train_dataset: DataParserConfig = FriendsDataParserConfig() image_dataset_type: str = "panoptic" -@dataclass -class InstantNGPDatasetConfig(DatasetConfig): - """Mipnerf 360 dataset config""" - - _target: Type = InstantNGP - data_directory: Path = Path("data/ours/posterv2") - scale_factor: float = 1.0 - downscale_factor: int = 1 - scene_scale: float = 0.33 - - # Model related configs @dataclass class ModelConfig(InstantiateConfig): @@ -311,7 +311,7 @@ class PipelineConfig(InstantiateConfig): """Configuration for pipeline instantiation""" _target: Type = Pipeline - dataloader: VanillaDataloaderConfig = VanillaDataloaderConfig() + data_manager: VanillaDataManagerConfig = VanillaDataManagerConfig() model: ModelConfig = ModelConfig() diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py index 09331f447..ea6eff6a3 100644 --- a/nerfactory/configs/base_configs.py +++ b/nerfactory/configs/base_configs.py @@ -19,11 +19,11 @@ Put all the method implementations in one location. from typing import Dict from nerfactory.configs.base import ( - BlenderDatasetConfig, + BlenderDataParserConfig, Config, - FriendsDataloaderConfig, + FriendsDataManagerConfig, InstantNGPModelConfig, - MipNerf360DatasetConfig, + MipNerf360DataParserConfig, ModelConfig, NerfWModelConfig, OptimizerConfig, @@ -31,7 +31,7 @@ from nerfactory.configs.base import ( SchedulerConfig, TensoRFModelConfig, TrainerConfig, - VanillaDataloaderConfig, + VanillaDataManagerConfig, ) from nerfactory.models.mipnerf import MipNerfModel from nerfactory.models.mipnerf_360 import MipNerf360Model @@ -43,8 +43,8 @@ base_configs["instant_ngp"] = Config( method_name="instant_ngp", trainer=TrainerConfig(mixed_precision=True), pipeline=PipelineConfig( - dataloader=VanillaDataloaderConfig( - train_dataset=BlenderDatasetConfig(), + data_manager=VanillaDataManagerConfig( + train_dataset=BlenderDataParserConfig(), train_num_rays_per_batch=8192, eval_num_rays_per_chunk=8192, ), @@ -63,8 +63,8 @@ base_configs["mipnerf_360"] = Config( method_name="mipnerf_360", trainer=TrainerConfig(steps_per_test=200), pipeline=PipelineConfig( - dataloader=VanillaDataloaderConfig( - train_dataset=MipNerf360DatasetConfig(), + data_manager=VanillaDataManagerConfig( + train_dataset=MipNerf360DataParserConfig(), train_num_rays_per_batch=8192, eval_num_rays_per_chunk=8192, ), @@ -81,8 +81,8 @@ base_configs["mipnerf_360"] = Config( base_configs["mipnerf"] = Config( method_name="mipnerf", pipeline=PipelineConfig( - dataloader=VanillaDataloaderConfig( - train_dataset=BlenderDatasetConfig(), + data_manager=VanillaDataManagerConfig( + train_dataset=BlenderDataParserConfig(), train_num_rays_per_batch=8192, eval_num_rays_per_chunk=8192, ), @@ -98,14 +98,14 @@ base_configs["mipnerf"] = Config( base_configs["nerfw"] = Config( experiment_name="friends_TBBT-big_living_room", method_name="nerfw", - pipeline=PipelineConfig(dataloader=FriendsDataloaderConfig(), model=NerfWModelConfig()), + pipeline=PipelineConfig(data_manager=FriendsDataManagerConfig(), model=NerfWModelConfig()), ) base_configs["semantic_nerf"] = Config( experiment_name="friends_TBBT-big_living_room", method_name="semantic_nerf", pipeline=PipelineConfig( - dataloader=FriendsDataloaderConfig(), + data_manager=FriendsDataManagerConfig(), model=ModelConfig( _target=SemanticNerfModel, loss_coefficients={"rgb_loss_coarse": 1.0, "rgb_loss_fine": 1.0, "semantic_loss_fine": 0.05}, @@ -118,8 +118,8 @@ base_configs["semantic_nerf"] = Config( base_configs["vanilla_nerf"] = Config( method_name="vanilla_nerf", pipeline=PipelineConfig( - dataloader=VanillaDataloaderConfig( - train_dataset=BlenderDatasetConfig(), + data_manager=VanillaDataManagerConfig( + train_dataset=BlenderDataParserConfig(), ), model=ModelConfig(_target=NeRFModel), ), @@ -129,8 +129,8 @@ base_configs["tensorf"] = Config( method_name="tensorf", trainer=TrainerConfig(mixed_precision=True), pipeline=PipelineConfig( - dataloader=VanillaDataloaderConfig( - train_dataset=BlenderDatasetConfig(), + data_manager=VanillaDataManagerConfig( + train_dataset=BlenderDataParserConfig(), ), model=TensoRFModelConfig(), ), diff --git a/nerfactory/dataloaders/base.py b/nerfactory/dataloaders/base.py index 7cf2d0857..fa2fb26d7 100644 --- a/nerfactory/dataloaders/base.py +++ b/nerfactory/dataloaders/base.py @@ -24,12 +24,16 @@ from typing import Dict, List, Optional, Tuple, Union import torch from torch import nn from torch.nn import Parameter +from torch.utils.data.distributed import DistributedSampler from nerfactory.cameras.rays import RayBundle from nerfactory.configs import base as cfg -from nerfactory.dataloaders.eval import EvalDataloader, FixedIndicesEvalDataloader +from nerfactory.dataloaders.image_dataloader import ( + CacheImageDataloader, + EvalDataloader, + FixedIndicesEvalDataloader, +) from nerfactory.dataloaders.image_dataset import ImageDataset, PanopticImageDataset -from nerfactory.dataloaders.image_sampler import CacheImageSampler from nerfactory.dataloaders.pixel_sampler import PixelSampler from nerfactory.dataloaders.structs import DatasetInputs from nerfactory.models.modules.ray_generator import RayGenerator @@ -37,17 +41,17 @@ from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttribu from nerfactory.utils.misc import IterableWrapper -class Dataloader(nn.Module): - """Generic dataloader's abstract class +class DataManager(nn.Module): + """Generic data manager's abstract class - This version of the dataloader is designed be a monolithic way to load data and latents, + This version of the data manager is designed be a monolithic way to load data and latents, especially since this may contain learnable parameters which need to be shared across the train - and test dataloaders. The idea is that we have setup methods for train and eval separately and + and test data managers. The idea is that we have setup methods for train and eval separately and this can be a combined train/eval if you want. Usage: To get data, use the next_train and next_eval functions. - This dataloader's next_train and next_eval methods will return 2 things: + This data manager's next_train and next_eval methods will return 2 things: 1. A Raybundle: This will contain the rays we are sampling, with latents and conditionals attached (everything needed at inference) 2. A "batch" of auxilury information: This will contain the mask, the ground truth @@ -72,8 +76,8 @@ class Dataloader(nn.Module): get_eval_iterable: utility that gets a clean pythonic iterator for your eval data Args: - use_train: whether this dataloader is being used for training - use_eval: whether this dataloader is being used for evaluation + use_train: whether this data manager is being used for training + use_eval: whether this data manager is being used for evaluation Attributes: use_train (bool): whether or not we are using train @@ -95,9 +99,9 @@ class Dataloader(nn.Module): eval_dataloader: Optional[EvalDataloader] = None def __init__(self, use_train: bool, use_eval: bool): - """Constructor for the Dataloader class. + """Constructor for the DataManager class. - Subclassed Dataloaders will likely need to override this constructor. + Subclassed DataManagers will likely need to override this constructor. If you aren't manually calling the setup_train and setup_eval functions from an overriden constructor, that you call super().__init__(use_train, use_eval) BEFORE you initialize any @@ -140,10 +144,10 @@ class Dataloader(nn.Module): as __iter__ and __next__ methods respectivley. This basically is just a little utility if you want to do something like: - | for ray_bundle, batch in dataloader.get_train_iterable(): + | for ray_bundle, batch in data_manager.get_train_iterable(): | since the returned IterableWrapper is just an iterator with the __iter__ and __next__ - methods (methods bound to our Dataloader instance in this case) specified in the constructor. + methods (methods bound to our DataManager instance in this case) specified in the constructor. """ return IterableWrapper(self.iter_train, self.next_train, length) @@ -152,38 +156,38 @@ class Dataloader(nn.Module): as __iter__ and __next__ methods respectivley. This basically is just a little utility if you want to do something like: - | for ray_bundle, batch in dataloader.get_eval_iterable(): + | for ray_bundle, batch in data_manager.get_eval_iterable(): | since the returned IterableWrapper is just an iterator with the __iter__ and __next__ - methods (methods bound to our Dataloader instance in this case) specified in the constructor. + methods (methods bound to our DataManager instance in this case) specified in the constructor. """ return IterableWrapper(self.iter_eval, self.next_eval, length) @abstractmethod def setup_train(self): - """Sets up the dataloader for training. + """Sets up the data manager for training. Here you will define any subclass specific object attributes from the attribute""" raise NotImplementedError @abstractmethod def setup_eval(self): - """Sets up the dataloader for evaluation""" + """Sets up the data manager for evaluation""" raise NotImplementedError @abstractmethod def next_train(self) -> Tuple: - """Returns the next batch of data from the train dataloader. + """Returns the next batch of data from the train data manager. - This will be a tuple of all the information that this dataloader outputs. + This will be a tuple of all the information that this data manager outputs. """ raise NotImplementedError @abstractmethod def next_eval(self) -> Tuple: - """Returns the next batch of data from the eval dataloader. + """Returns the next batch of data from the eval data manager. - This will be a tuple of all the information that this dataloader outputs. + This will be a tuple of all the information that this data manager outputs. """ raise NotImplementedError @@ -195,16 +199,16 @@ class Dataloader(nn.Module): @abstractmethod def get_param_groups(self) -> Dict[str, List[Parameter]]: # pylint: disable=no-self-use - """Get the param groups for the dataloader. + """Get the param groups for the data manager. Returns: - A list of dictionaries containing the dataloader's param groups. + A list of dictionaries containing the data manager's param groups. """ return {} -class VanillaDataloader(Dataloader): # pylint: disable=abstract-method - """Basic stored dataloader implementation. +class VanillaDataManager(DataManager): # pylint: disable=abstract-method + """Basic stored data manager implementation. This is pretty much a port over from our old dataloading utilities, and is a little jank under the hood. We may clean this up a little bit under the hood with more standard dataloading @@ -212,17 +216,22 @@ class VanillaDataloader(Dataloader): # pylint: disable=abstract-method only the constructor is likely to change in the future, or maybe passing in step number to the next_train and next_eval functions.""" - config: cfg.VanillaDataloaderConfig + config: cfg.VanillaDataManagerConfig def __init__( self, - config: cfg.VanillaDataloaderConfig, + config: cfg.VanillaDataManagerConfig, device: Union[torch.device, str] = "cpu", test_mode: bool = False, + world_size: int = 1, + local_rank: int = 0, **kwargs, # pylint: disable=unused-argument ): self.config = config self.device = device + self.world_size = world_size + self.local_rank = local_rank + self.sampler = None dataset_train = config.train_dataset.setup() self.train_datasetinputs = dataset_train.get_dataset_inputs(split="train") @@ -239,43 +248,68 @@ class VanillaDataloader(Dataloader): # pylint: disable=abstract-method super().__init__(use_train, use_eval) def setup_train(self): - """Sets up the dataloader for training""" + """Sets up the data manager for training""" assert self.train_datasetinputs is not None if self.config.image_dataset_type == "rgb": - self.train_image_dataset = ImageDataset(**self.train_datasetinputs.as_dict()) + self.train_image_dataset = ImageDataset( + image_filenames=self.train_datasetinputs.image_filenames, + downscale_factor=self.train_datasetinputs.downscale_factor, + alpha_color=self.train_datasetinputs.alpha_color, + ) elif self.config.image_dataset_type == "panoptic": self.train_image_dataset = PanopticImageDataset(**self.train_datasetinputs.as_dict()) else: raise ValueError(f"Unknown image dataset type {self.config.image_dataset_type}") - self.train_image_sampler = CacheImageSampler( - self.train_image_dataset, - num_images_to_sample_from=self.config.train_num_images_to_sample_from, - device=self.device, - ) # TODO(ethan): pass this in - self.iter_train_image_sampler = iter(self.train_image_sampler) + + if self.world_size > 1: + self.train_image_dataloader = CacheImageDataloader( + self.train_image_dataset, + num_images_to_sample_from=self.config.train_num_images_to_sample_from, + device=self.device, + num_workers=self.world_size * 4, + pin_memory=True, + ) # TODO(ethan): pass this in + else: + self.train_image_dataloader = CacheImageDataloader( + self.train_image_dataset, + num_images_to_sample_from=self.config.train_num_images_to_sample_from, + device=self.device, + ) + self.iter_train_image_dataloader = iter(self.train_image_dataloader) self.train_pixel_sampler = PixelSampler(self.config.train_num_rays_per_batch) self.train_ray_generator = RayGenerator(self.train_cameras) def setup_eval(self): - """Sets up the dataloader for evaluation""" + """Sets up the data manager for evaluation""" assert self.eval_datasetinputs is not None if self.config.image_dataset_type == "rgb": - self.eval_image_dataset = ImageDataset(**self.eval_datasetinputs.as_dict()) + self.eval_image_dataset = ImageDataset( + image_filenames=self.eval_datasetinputs.image_filenames, + downscale_factor=self.eval_datasetinputs.downscale_factor, + alpha_color=self.eval_datasetinputs.alpha_color, + ) elif self.config.image_dataset_type == "panoptic": self.eval_image_dataset = PanopticImageDataset(**self.eval_datasetinputs.as_dict()) + + if self.world_size > 1: + self.sampler = DistributedSampler( + self.train_image_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42 + ) + self.eval_dataloader = FixedIndicesEvalDataloader( image_dataset=self.eval_image_dataset, cameras=self.eval_cameras, num_rays_per_chunk=self.config.eval_num_rays_per_chunk, image_indices=self.config.eval_image_indices, device=self.device, + num_workers=0 if self.world_size == 1 else self.world_size * 4, ) def next_train(self) -> Tuple[RayBundle, Dict]: """Returns the next batch of data from the train dataloader. The RayBundle can be shaped in whatever way.""" self.train_count += 1 - image_batch = next(self.iter_train_image_sampler) + image_batch = next(self.iter_train_image_dataloader) batch = self.train_pixel_sampler.sample(image_batch) ray_indices = batch["indices"] ray_bundle = self.train_ray_generator(ray_indices) diff --git a/nerfactory/dataloaders/datasets.py b/nerfactory/dataloaders/data_parsers.py similarity index 98% rename from nerfactory/dataloaders/datasets.py rename to nerfactory/dataloaders/data_parsers.py index e1d743347..4b33432ae 100644 --- a/nerfactory/dataloaders/datasets.py +++ b/nerfactory/dataloaders/data_parsers.py @@ -52,10 +52,10 @@ from nerfactory.utils.misc import get_hash_str_from_dict @dataclass -class Dataset: +class DataParser: """A dataset.""" - def __init__(self, config: cfg.DatasetConfig): + def __init__(self, config: cfg.DataParserConfig): super().__init__() self.config = config @@ -130,7 +130,7 @@ class Dataset: @dataclass -class Blender(Dataset): +class Blender(DataParser): """Blender Dataset Some of this code comes from https://github.com/yenchenlin/nerf-pytorch/blob/master/load_blender.py#L37. @@ -141,7 +141,7 @@ class Blender(Dataset): downscale_factor: How much to downscale images. Defaults to 1. """ - def __init__(self, config: cfg.BlenderDatasetConfig): + def __init__(self, config: cfg.BlenderDataParserConfig): super().__init__(config=config) self.data_directory: Path = config.data_directory self.scale_factor: float = config.scale_factor @@ -200,7 +200,7 @@ class Blender(Dataset): @dataclass -class InstantNGP(Dataset): +class InstantNGP(DataParser): """Instant NGP Dataset Args: @@ -210,7 +210,7 @@ class InstantNGP(Dataset): scene_scale: How much to scale the scene. Defaults to 0.33 """ - config: cfg.InstantNGPDatasetConfig + config: cfg.InstantNGPDataParserConfig def _generate_dataset_inputs(self, split="train"): @@ -288,7 +288,7 @@ class InstantNGP(Dataset): @dataclass -class Mipnerf360(Dataset): +class Mipnerf360(DataParser): """MipNeRF 360 Dataset Args: @@ -299,7 +299,7 @@ class Mipnerf360(Dataset): aabb_scale: Scene scale, Defaults to 1.0. """ - config: cfg.MipNerf360DatasetConfig + config: cfg.MipNerf360DataParserConfig @classmethod def normalize_orientation(cls, poses: np.ndarray): @@ -412,7 +412,7 @@ class Mipnerf360(Dataset): @dataclass -class Record3D(Dataset): +class Record3D(DataParser): """Record3D Dataset Args: @@ -424,7 +424,7 @@ class Record3D(Dataset): more, images will be sampled approximately evenly. Defaults to 150. """ - config: cfg.Record3DDatasetConfig + config: cfg.Record3DDataParserConfig def _generate_dataset_inputs(self, split: str = "train") -> DatasetInputs: abs_dir = get_absolute_path(self.config.data_directory) @@ -534,7 +534,7 @@ class Record3D(Dataset): @dataclass -class Friends(Dataset): +class Friends(DataParser): """Friends Dataset Args: diff --git a/nerfactory/dataloaders/eval.py b/nerfactory/dataloaders/image_dataloader.py similarity index 64% rename from nerfactory/dataloaders/eval.py rename to nerfactory/dataloaders/image_dataloader.py index e57b46888..a90d26ad0 100644 --- a/nerfactory/dataloaders/eval.py +++ b/nerfactory/dataloaders/image_dataloader.py @@ -13,7 +13,7 @@ # limitations under the License. """ -Data loader. +Code for sampling images from a dataset of images. """ import random @@ -21,6 +21,8 @@ from abc import abstractmethod from typing import Dict, Optional, Tuple, Union import torch +from torch.utils.data import default_collate +from torch.utils.data.dataloader import DataLoader from nerfactory.cameras.cameras import Cameras from nerfactory.cameras.rays import RayBundle @@ -28,7 +30,76 @@ from nerfactory.dataloaders.image_dataset import ImageDataset from nerfactory.utils.misc import get_dict_to_torch -class EvalDataloader: +class CacheImageDataloader(DataLoader): + """Collated image dataset that implements caching of images. + Creates batches of the ImageDataset return type. + + Args: + dataset (torch.utils.data.Dataset): Dataset to sample from. + num_samples_to_collate (int, optional): How many images to sample rays for each batch. Defaults to all images. + num_times_to_repeat_images (int): How often to collate new images. Defaults to every iteration. + device (Union[torch.device, str]): Device to perform computation. Defaults to "cpu". + """ + + def __init__( + self, + dataset: ImageDataset, + num_images_to_sample_from: int = -1, + num_times_to_repeat_images: int = 0, + device: Union[torch.device, str] = "cpu", + **kwargs, + ): + self.dataset = dataset + self.num_times_to_repeat_images = num_times_to_repeat_images + self.cache_all_images = num_images_to_sample_from == -1 + self.num_images_to_sample_from = len(self.dataset) if self.cache_all_images else num_images_to_sample_from + self.device = device + + self.num_repeated = self.num_times_to_repeat_images # starting value + self.first_time = True + + self.cached_collated_batch = None + if self.cache_all_images: + self.cached_collated_batch = self._get_collated_batch() + super().__init__(dataset=dataset, **kwargs) + + def __getitem__(self, idx): + return self.dataset.__getitem__(idx) + + def _get_batch_list(self): + """Returns a list of batches from the dataset attribute.""" + indices = random.sample(range(len(self.dataset)), k=self.num_images_to_sample_from) + + batch_list = [self.dataset.__getitem__(idx) for idx in indices] + return batch_list + + def _get_collated_batch(self): + """Returns a collated batch.""" + batch_list = self._get_batch_list() + collated_batch = default_collate(batch_list) + collated_batch = get_dict_to_torch(collated_batch, device=self.device) + return collated_batch + + def __iter__(self): + while True: + if self.cache_all_images: + collated_batch = self.cached_collated_batch + elif self.first_time or ( + self.num_times_to_repeat_images != -1 and self.num_repeated >= self.num_times_to_repeat_images + ): + # trigger a reset + self.num_repeated = 0 + collated_batch = self._get_collated_batch() + # possibly save a cached item + self.cached_collated_batch = collated_batch if self.num_times_to_repeat_images != 0 else None + self.first_time = False + else: + collated_batch = self.cached_collated_batch + self.num_repeated += 1 + yield collated_batch + + +class EvalDataloader(DataLoader): """Evaluation dataloader base class Args: @@ -46,12 +117,12 @@ class EvalDataloader: device: Union[torch.device, str] = "cpu", **kwargs, ): - super().__init__() self.image_dataset = image_dataset self.cameras = cameras self.num_rays_per_chunk = num_rays_per_chunk self.device = device self.kwargs = kwargs + super().__init__(dataset=image_dataset) @abstractmethod def __iter__(self): diff --git a/nerfactory/dataloaders/image_dataset.py b/nerfactory/dataloaders/image_dataset.py index 95f04b58b..6e8fb1676 100644 --- a/nerfactory/dataloaders/image_dataset.py +++ b/nerfactory/dataloaders/image_dataset.py @@ -39,7 +39,6 @@ class ImageDataset(Dataset): image_filenames: List[Path], downscale_factor: int = 1, alpha_color: Optional[TensorType[3]] = None, - **kwargs, ): """_summary_ @@ -53,7 +52,6 @@ class ImageDataset(Dataset): self.image_filenames = image_filenames self.downscale_factor = downscale_factor self.alpha_color = alpha_color - self.kwargs = kwargs def __len__(self): return len(self.image_filenames) diff --git a/nerfactory/dataloaders/image_sampler.py b/nerfactory/dataloaders/image_sampler.py deleted file mode 100644 index 436813f96..000000000 --- a/nerfactory/dataloaders/image_sampler.py +++ /dev/null @@ -1,116 +0,0 @@ -# Copyright 2022 The Plenoptix Team. 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. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Code for sampling images from a dataset of images. -""" - -import random -from abc import abstractmethod -from typing import Union - -import torch -from torch.utils.data import IterableDataset, default_collate - -from nerfactory.dataloaders.image_dataset import ImageDataset -from nerfactory.utils.misc import get_dict_to_torch - - -class ImageSampler(IterableDataset): - """Samples image_batch's. - - Attributes: - dataset (torch.utils.data.Dataset): Dataset to sample from. - """ - - # dataset: ImageDataset - def __init__(self, dataset: ImageDataset): - super().__init__() - self.dataset = dataset - - @abstractmethod - def __getitem__(self, idx): - raise NotImplementedError() - - @abstractmethod - def __iter__(self): - raise NotImplementedError() - - -class CacheImageSampler(ImageSampler): - """Collated image dataset that implements caching of images. - Creates batches of the ImageDataset return type. - - Args: - dataset (torch.utils.data.Dataset): Dataset to sample from. - num_samples_to_collate (int, optional): How many images to sample rays for each batch. Defaults to all images. - num_times_to_repeat_images (int): How often to collate new images. Defaults to every iteration. - device (Union[torch.device, str]): Device to perform computation. Defaults to "cpu". - """ - - def __init__( - self, - dataset: ImageDataset, - num_images_to_sample_from: int = -1, - num_times_to_repeat_images: int = 0, - device: Union[torch.device, str] = "cpu", - ): - super().__init__(dataset=dataset) - self.dataset = dataset - self.cache_all_images = num_images_to_sample_from == -1 - self.num_images_to_sample_from = len(self.dataset) if self.cache_all_images else num_images_to_sample_from - self.num_times_to_repeat_images = num_times_to_repeat_images - self.device = device - - self.num_repeated = self.num_times_to_repeat_images # starting value - self.first_time = True - - self.cached_collated_batch = None - if self.cache_all_images: - self.cached_collated_batch = self._get_collated_batch() - - def __getitem__(self, idx): - return self.dataset.__getitem__(idx) - - def _get_batch_list(self): - """Returns a list of batches from the dataset attribute.""" - indices = random.sample(range(len(self.dataset)), k=self.num_images_to_sample_from) - - batch_list = [self.dataset.__getitem__(idx) for idx in indices] - return batch_list - - def _get_collated_batch(self): - """Returns a collated batch.""" - batch_list = self._get_batch_list() - collated_batch = default_collate(batch_list) - collated_batch = get_dict_to_torch(collated_batch, device=self.device) - return collated_batch - - def __iter__(self): - while True: - if self.cache_all_images: - collated_batch = self.cached_collated_batch - elif self.first_time or ( - self.num_times_to_repeat_images != -1 and self.num_repeated >= self.num_times_to_repeat_images - ): - # trigger a reset - self.num_repeated = 0 - collated_batch = self._get_collated_batch() - # possibly save a cached item - self.cached_collated_batch = collated_batch if self.num_times_to_repeat_images != 0 else None - self.first_time = False - else: - collated_batch = self.cached_collated_batch - self.num_repeated += 1 - yield collated_batch diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 6cb56d236..458ae51b8 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -19,14 +19,11 @@ from __future__ import annotations import functools import logging -import typing from pathlib import Path from typing import Any, Dict, List import torch -import torch.distributed as dist from torch.cuda.amp.grad_scaler import GradScaler -from torch.nn.parallel import DistributedDataParallel as DDP from nerfactory.configs import base as cfg from nerfactory.optimizers.optimizers import Optimizers, setup_optimizers @@ -103,29 +100,26 @@ class Trainer: Args: test_mode (bool, optional): Whether to setup for testing. Defaults to False. """ - self.pipeline: Pipeline = self.config.pipeline.setup(device=self.device, test_mode=test_mode) + self.pipeline: Pipeline = 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()) self._load_checkpoint() - if self.world_size > 1: - self.pipeline = typing.cast( - Pipeline, typing.cast(Pipeline, DDP(self.pipeline, device_ids=[self.local_rank])) - ) - dist.barrier(device_ids=[self.local_rank]) - + # TODO(ethan): do this for pipeline, not pipeline.model self.callbacks = self.pipeline.get_training_callbacks( TrainingCallbackAttributes(optimizers=self.optimizers, grad_scaler=self.grad_scaler, pipeline=self.pipeline) ) def train(self) -> None: """Train the model.""" - assert self.pipeline.dataloader.train_datasetinputs is not None, "Missing DatasetInputs" - assert self.pipeline.dataloader.train_image_dataset is not None, "Missing ImageDataset" + assert self.pipeline.data_manager.train_datasetinputs is not None, "Missing DatasetInputs" + assert self.pipeline.data_manager.train_image_dataset is not None, "Missing ImageDataset" self.visualizer_state.init_scene( - image_dataset=self.pipeline.dataloader.train_image_dataset, - dataset_inputs=self.pipeline.dataloader.train_datasetinputs, + image_dataset=self.pipeline.data_manager.train_image_dataset, + dataset_inputs=self.pipeline.data_manager.train_datasetinputs, ) with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME): num_iterations = self.config.trainer.max_num_iterations @@ -149,7 +143,7 @@ class Trainer: with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _: self.visualizer_state.update_scene(step, self.pipeline.model) - train_num_rays_per_batch = self.config.pipeline.dataloader.train_num_rays_per_batch + train_num_rays_per_batch = self.config.pipeline.data_manager.train_num_rays_per_batch writer.put_scalar(name=EventName.RAYS_PER_SEC, scalar=train_num_rays_per_batch / t.duration, step=step) if step != 0 and step % self.config.logging.steps_per_log == 0: diff --git a/nerfactory/models/base.py b/nerfactory/models/base.py index a502738c8..402c896e7 100644 --- a/nerfactory/models/base.py +++ b/nerfactory/models/base.py @@ -144,6 +144,7 @@ class Model(nn.Module): This outputs different things depending on the configuration of the model and whether or not the batch is provided (whether or not we are training basically).""" + if self.collider is not None: intersected_ray_bundle = self.collider(ray_bundle) # pylint: disable=not-callable valid_mask = intersected_ray_bundle.valid_mask[..., 0] diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index 6e3c82a30..9b1dac099 100644 --- a/nerfactory/pipelines/base.py +++ b/nerfactory/pipelines/base.py @@ -17,14 +17,17 @@ Abstracts for the Pipeline class. """ from __future__ import annotations +import typing from abc import abstractmethod from typing import Any, Dict, List, Optional +import torch.distributed as dist from torch import nn from torch.nn import Parameter +from torch.nn.parallel import DistributedDataParallel as DDP from nerfactory.configs import base as cfg -from nerfactory.dataloaders.base import Dataloader +from nerfactory.dataloaders.base import DataManager from nerfactory.models.base import Model from nerfactory.utils import profiler from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes @@ -40,9 +43,9 @@ class Pipeline(nn.Module): class should be 1:1 with a pipeline that can act as a standardized interface and hide differences in how each model takes in and outputs data. - This class's function is to hide the dataloader and model classes from the trainer, + This class's function is to hide the data manager and model classes from the trainer, worrying about: - 1) Fetching data with the dataloader + 1) Fetching data with the data manager 2) Feeding the model the data and fetching the loss Hopefully this provides a higher level interface for the trainer to use, and simplifying the model classes, which each may have different forward() methods @@ -51,47 +54,60 @@ class Pipeline(nn.Module): TODO: For visualizer functionality to be added down the line, we should make sure that we are still abstracting away the Model from the end user. The visualizer function - should probably be done by adding a new iterator on the base dataloader that will + should probably be done by adding a new iterator on the base data manager that will talk to the actual visualizer. This is probably ideal to have the visualizer be - primarily located in the dataloader (first because it makes sense as the + primarily located in the data manager (first because it makes sense as the visualizers 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 visualizer 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 dataloader, + be done in the future... but just bear in mind that if learned parameters are in the data manager, the visualizer may have to use those parameters as well. Args: model: The model to be used in the pipeline. - dataloader: The dataloader to be used in the pipeline. + data_manager: The data_manager to be used in the pipeline. loss_coefficients: A dictionary of loss coefficients that will be used Attributes: - self.dataloader (Dataloader): The dataloader that will be used + self.data_manager (DataManager): The data manager that will be used self.model (Model): The model that will be used """ - def __init__(self, config: cfg.PipelineConfig, device: str, test_mode: bool = False): + def __init__( + self, config: cfg.PipelineConfig, device: str, test_mode: bool = False, world_size: int = 1, local_rank: int = 0 + ): super().__init__() - self.dataloader: Dataloader = config.dataloader.setup(device=device, test_mode=test_mode) - self.dataloader.to(device) + self.data_manager: DataManager = config.data_manager.setup( + device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank + ) + self.data_manager.to(device) # TODO(ethan): get rid of scene_bounds from the model - assert self.dataloader.train_datasetinputs is not None, "Missing DatasetInputs" - self.model: Model = config.model.setup(scene_bounds=self.dataloader.train_datasetinputs.scene_bounds) + assert self.data_manager.train_datasetinputs is not None, "Missing DatasetInputs" + self.model: Model = config.model.setup(scene_bounds=self.data_manager.train_datasetinputs.scene_bounds) self.model.to(device) + self.world_size = world_size + if world_size > 1: + self.model = typing.cast( + Model, typing.cast(Model, DDP(self.model, device_ids=[local_rank], find_unused_parameters=True)) + ) + dist.barrier(device_ids=[local_rank]) + @property def device(self): """Returns the device that the model is on.""" - return self.model.device + return self.model.module.device if self.world_size > 1 else self.model.device @abstractmethod @profiler.time_function def get_train_loss_dict(self, step: Optional[int] = None): """This function gets your training loss dict. This will be responsible for - getting the next batch of data from the dataloader and interfacing with the + getting the next batch of data from the DataManager and interfacing with the Model class, feeding the data to the model's forward function.""" - ray_bundle, batch = self.dataloader.next_train() + if self.world_size > 1: + self.data_manager.sampler.set_epoch(step) + ray_bundle, batch = self.data_manager.next_train() model_outputs, loss_dict, metrics_dict = self.model(ray_bundle, batch) return model_outputs, loss_dict, metrics_dict @@ -99,15 +115,19 @@ class Pipeline(nn.Module): @profiler.time_function def get_eval_loss_dict(self, step: Optional[int] = None): """This function gets your evaluation loss dict. It needs to get the data - from the dataloader and feed it to the model's forward function""" + from the DataManager and feed it to the model's forward function""" self.eval() # NOTE(ethan): next_eval() is not being used right now - assert self.dataloader.eval_dataloader is not None - for camera_ray_bundle, batch in self.dataloader.eval_dataloader: + assert self.data_manager.eval_dataloader is not None + for camera_ray_bundle, batch in self.data_manager.eval_dataloader: assert camera_ray_bundle.camera_indices is not None image_idx = int(camera_ray_bundle.camera_indices[0, 0]) - outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - psnr = self.model.log_test_image_outputs(image_idx, step, batch, outputs) + if self.world_size > 1: + outputs = self.model.module.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + psnr = self.model.module.log_test_image_outputs(image_idx, step, batch, outputs) + else: + outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + psnr = self.model.log_test_image_outputs(image_idx, step, batch, outputs) # TODO(ethan): this function should probably return something? self.train() @@ -124,9 +144,9 @@ class Pipeline(nn.Module): self, training_callback_attributes: TrainingCallbackAttributes ) -> List[TrainingCallback]: """Returns the training callbacks from both the Dataloader and the Model.""" - dataloader_callbacks = self.dataloader.get_training_callbacks(training_callback_attributes) + data_manager_callbacks = self.data_manager.get_training_callbacks(training_callback_attributes) model_callbacks = self.model.get_training_callbacks(training_callback_attributes) - callbacks = dataloader_callbacks + model_callbacks + callbacks = data_manager_callbacks + model_callbacks return callbacks def get_param_groups(self) -> Dict[str, List[Parameter]]: @@ -135,7 +155,10 @@ class Pipeline(nn.Module): Returns: A list of dictionaries containing the pipeline's param groups. """ - dataloader_params = self.dataloader.get_param_groups() - model_params = self.model.get_param_groups() + data_manager_params = self.data_manager.get_param_groups() + if self.world_size > 1: + model_params = self.model.module.get_param_groups() + else: + model_params = self.model.get_param_groups() # TODO(ethan): assert that key names don't overlap - return {**dataloader_params, **model_params} + return {**data_manager_params, **model_params} diff --git a/scripts/benchmarking/launch_train_blender.sh b/scripts/benchmarking/launch_train_blender.sh index 80c91018d..a28373d79 100755 --- a/scripts/benchmarking/launch_train_blender.sh +++ b/scripts/benchmarking/launch_train_blender.sh @@ -55,7 +55,7 @@ for dataset in ${DATASETS[@]}; do python scripts/run_train.py \ ${config_name} \ --logging.writer.2.no-enable \ - --pipeline.dataloader.train-dataset.data-directory=data/blender/${dataset} \ + --pipeline.data_manager.train-dataset.data-directory=data/blender/${dataset} \ --experiment-name=blender_${dataset}_${tag} \ --trainer.model-dir=nerfactory_models/ \ --trainer.steps-per-save=16000 \ diff --git a/scripts/run_eval.py b/scripts/run_eval.py index 7b1bb2f19..bb59b4fbc 100644 --- a/scripts/run_eval.py +++ b/scripts/run_eval.py @@ -68,7 +68,7 @@ def _load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: def _render_stats_dict(pipeline: Pipeline) -> Dict[str, float]: - """Helper function to evaluate the pipeline on a dataloader. + """Helper function to evaluate the pipeline on a DataManager. Args: pipeline (Pipeline): Pipeline to evaluate @@ -79,7 +79,7 @@ def _render_stats_dict(pipeline: Pipeline) -> Dict[str, float]: avg_psnr = 0 avg_rays_per_sec = 0 avg_fps = 0 - for step, (camera_ray_bundle, batch) in tqdm(enumerate(pipeline.dataloader.eval_dataloader)): + for step, (camera_ray_bundle, batch) in tqdm(enumerate(pipeline.data_manager.eval_dataloader)): with TimeWriter(writer=None, name=None, write=False) as t: with torch.no_grad(): image_idx = int(camera_ray_bundle.camera_indices[0, 0]) @@ -142,9 +142,9 @@ def _eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, Path]: # load checkpoints from wherever they were saved # TODO: expose the ability to choose an arbitrary checkpoint config.trainer.load_dir = config.trainer.model_dir - config.pipeline.dataloader.eval_image_indices = None + config.pipeline.data_manager.eval_image_indices = None - # setup pipeline (which includes the dataloaders) + # setup pipeline (which includes the DataManager) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") pipeline = config.pipeline.setup(device=device, test_mode=True) assert isinstance(pipeline, Pipeline) @@ -210,11 +210,11 @@ class RenderTrajectory: # TODO(ethan): use camera information from parsing args if self.traj == "spiral": - camera_start = pipeline.dataloader.eval_dataloader.get_camera(image_idx=0) + camera_start = pipeline.data_manager.eval_dataloader.get_camera(image_idx=0) # TODO(ethan): pass in the up direction of the camera camera_path = get_spiral_path(camera_start, steps=30, radius=0.1) elif self.traj == "interp": - cameras = pipeline.dataloader.eval_dataloader.get_camera(image_idx=[0, 10]) + cameras = pipeline.data_manager.eval_dataloader.get_camera(image_idx=[0, 10]) camera_path = get_interpolated_camera_path(cameras, steps=30) else: assert_never(self.traj) @@ -225,7 +225,7 @@ class RenderTrajectory: output_filename=self.output_path, rendered_output_name=self.rendered_output_name, rendered_resolution_scaling_factor=self.rendered_resolution_scaling_factor, - num_rays_per_chunk=config.pipeline.dataloader.eval_num_rays_per_chunk, + num_rays_per_chunk=config.pipeline.data_manager.eval_num_rays_per_chunk, ) diff --git a/scripts/run_train.py b/scripts/run_train.py index f474ba65c..4e1c89b4b 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -164,7 +164,6 @@ def launch( config, timeout, ), - join=False, ) assert process_context is not None try: diff --git a/tests/test_run_train.py b/tests/test_run_train.py index b8a8a7836..8fc7a172b 100644 --- a/tests/test_run_train.py +++ b/tests/test_run_train.py @@ -22,15 +22,15 @@ def set_reduced_config(config: cfg.Config): config.trainer.max_num_iterations = 2 # reduce dataset factors; set dataset to test # switch to using the vanilla ImageDataset class - config.pipeline.dataloader.image_dataset_type = "rgb" + config.pipeline.data_manager.image_dataset_type = "rgb" - config.pipeline.dataloader.train_dataset = cfg.BlenderDatasetConfig( - data_directory="tests/data/lego_test", downscale_factor=16 + config.pipeline.data_manager.train_dataset = cfg.BlenderDataParserConfig( + data_directory=Path("tests/data/lego_test"), downscale_factor=16 ) - config.pipeline.dataloader.train_num_images_to_sample_from = 1 - config.pipeline.dataloader.train_num_rays_per_batch = 4 - config.pipeline.dataloader.eval_dataset = cfg.BlenderDatasetConfig( - data_directory="tests/data/lego_test", downscale_factor=16 + config.pipeline.data_manager.train_num_images_to_sample_from = 1 + config.pipeline.data_manager.train_num_rays_per_batch = 4 + config.pipeline.data_manager.eval_dataset = cfg.BlenderDataParserConfig( + data_directory=Path("tests/data/lego_test"), downscale_factor=16 ) # reduce model factors