diff --git a/nerfstudio/data/datamanagers.py b/nerfstudio/data/datamanagers.py index c3cf367c3..774808bde 100644 --- a/nerfstudio/data/datamanagers.py +++ b/nerfstudio/data/datamanagers.py @@ -15,6 +15,7 @@ """ Data loader. """ + from __future__ import annotations from abc import abstractmethod @@ -23,6 +24,7 @@ from typing import Dict, List, Optional, Tuple, Type, Union import torch import tyro +from rich.progress import Console from torch import nn from torch.nn import Parameter from torch.utils.data import Dataset @@ -49,6 +51,8 @@ from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttrib from nerfstudio.model_components.ray_generators import RayGenerator from nerfstudio.utils.misc import IterableWrapper +CONSOLE = Console(width=120) + AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of flags in subcommands. tyro.extras.subcommand_type_from_defaults( { @@ -242,10 +246,16 @@ class VanillaDataManagerConfig(InstantiateConfig): """Number of rays per batch to use per training iteration.""" train_num_images_to_sample_from: int = -1 """Number of images to sample during training iteration.""" + train_num_times_to_repeat_images: int = -1 + """When not training on all images, number of iterations before picking new + images. If -1, never pick new images.""" eval_num_rays_per_batch: int = 1024 """Number of rays per batch to use per eval iteration.""" eval_num_images_to_sample_from: int = -1 """Number of images to sample during eval iteration.""" + eval_num_times_to_repeat_images: int = -1 + """When not evaluating on all images, number of iterations before picking + new images. If -1, never pick new images.""" eval_image_indices: Optional[Tuple[int, ...]] = (0,) """Specifies the image indices to use during eval; if None, uses all.""" camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig() @@ -294,9 +304,11 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method def setup_train(self): """Sets up the data loaders for training""" assert self.train_dataset is not None + CONSOLE.print("Setting up training dataset...") self.train_image_dataloader = CacheDataloader( self.train_dataset, num_images_to_sample_from=self.config.train_num_images_to_sample_from, + num_times_to_repeat_images=self.config.train_num_times_to_repeat_images, device=self.device, num_workers=self.world_size * 4, pin_memory=True, @@ -314,9 +326,11 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method def setup_eval(self): """Sets up the data loader for evaluation""" assert self.eval_dataset is not None + CONSOLE.print("Setting up evaluation dataset...") self.eval_image_dataloader = CacheDataloader( self.eval_dataset, num_images_to_sample_from=self.config.eval_num_images_to_sample_from, + num_times_to_repeat_images=self.config.eval_num_times_to_repeat_images, device=self.device, num_workers=self.world_size * 4, pin_memory=True, @@ -340,9 +354,6 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method num_workers=self.world_size * 4, ) - # TODO: eval dataloader should be separate from train - self.iter_eval_dataloader = iter(self.eval_image_dataloader) - def next_train(self, step: int) -> Tuple[RayBundle, Dict]: """Returns the next batch of data from the train dataloader.""" self.train_count += 1 diff --git a/nerfstudio/data/dataparsers/record3d_dataparser.py b/nerfstudio/data/dataparsers/record3d_dataparser.py index 513f9b4c6..cd3a94039 100644 --- a/nerfstudio/data/dataparsers/record3d_dataparser.py +++ b/nerfstudio/data/dataparsers/record3d_dataparser.py @@ -51,8 +51,8 @@ class Record3DDataParserConfig(DataParserConfig): orientation_method: Literal["pca", "up"] = "up" """The method to use for orientation""" max_dataset_size: int = 300 - """Max number of images to train on. If the dataset has - more, images will be sampled approximately evenly.""" + """Max number of images to train on. If the dataset has more, images will be sampled approximately evenly. If -1, + use all images.""" @dataclass @@ -86,7 +86,7 @@ class Record3D(DataParser): axis=-1, ).astype(np.float32) - if num_images > self.config.max_dataset_size: + if self.config.max_dataset_size != -1 and num_images > self.config.max_dataset_size: # Evenly select max_dataset_size images from dataset, including first # and last indices. idx = np.round(np.linspace(0, num_images - 1, self.config.max_dataset_size)).astype(int) diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index fab0db183..825cb65a1 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/nerfstudio/data/utils/dataloaders.py @@ -21,7 +21,7 @@ from abc import abstractmethod from typing import Dict, Optional, Tuple, Union import torch -from rich.progress import track +from rich.progress import Console, track from torch.utils.data import Dataset, default_collate from torch.utils.data.dataloader import DataLoader @@ -30,6 +30,8 @@ from nerfstudio.cameras.rays import RayBundle from nerfstudio.data.utils.datasets import InputDataset from nerfstudio.utils.misc import get_dict_to_torch +CONSOLE = Console(width=120) + class CacheDataloader(DataLoader): """Collated image dataset that implements caching of default-pytorch-collatable data. @@ -38,7 +40,7 @@ class CacheDataloader(DataLoader): Args: dataset: Dataset to sample from. num_samples_to_collate: How many images to sample rays for each batch. -1 for all images. - num_times_to_repeat_images: How often to collate new images. + num_times_to_repeat_images: How often to collate new images. -1 to never pick new images. device: Device to perform computation. """ @@ -46,13 +48,13 @@ class CacheDataloader(DataLoader): self, dataset: Dataset, num_images_to_sample_from: int = -1, - num_times_to_repeat_images: int = 0, + num_times_to_repeat_images: int = -1, 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.cache_all_images = (num_images_to_sample_from == -1) or (num_images_to_sample_from >= len(self.dataset)) self.num_images_to_sample_from = len(self.dataset) if self.cache_all_images else num_images_to_sample_from self.device = device @@ -61,7 +63,21 @@ class CacheDataloader(DataLoader): self.cached_collated_batch = None if self.cache_all_images: + CONSOLE.print(f"Caching all {len(self.dataset)} images.") + if len(self.dataset) > 500: + CONSOLE.print( + "[bold yellow]Warning: If you run out of memory, try reducing the number of images to sample from." + ) self.cached_collated_batch = self._get_collated_batch() + elif self.num_times_to_repeat_images == -1: + CONSOLE.print( + f"Caching {self.num_images_to_sample_from} out of {len(self.dataset)} images, without resampling." + ) + else: + CONSOLE.print( + f"Caching {self.num_images_to_sample_from} out of {len(self.dataset)} images, " + f"resampling every {self.num_times_to_repeat_images} iters." + ) super().__init__(dataset=dataset, **kwargs) def __getitem__(self, idx):