From 3489e992b0d497252e3061b22f899e39322f9184 Mon Sep 17 00:00:00 2001 From: Jake Austin <55679937+jake-austin@users.noreply.github.com> Date: Sun, 2 Oct 2022 12:24:05 -0700 Subject: [PATCH] Datamanager fields typing (#609) --- nerfstudio/data/datamanagers.py | 39 +++++++++++++++++-------------- nerfstudio/data/utils/__init__.py | 13 +++++++++++ nerfstudio/engine/trainer.py | 6 ++--- nerfstudio/pipelines/base.py | 6 ++--- 4 files changed, 40 insertions(+), 24 deletions(-) create mode 100644 nerfstudio/data/utils/__init__.py diff --git a/nerfstudio/data/datamanagers.py b/nerfstudio/data/datamanagers.py index 9389dfd17..3ecd96b04 100644 --- a/nerfstudio/data/datamanagers.py +++ b/nerfstudio/data/datamanagers.py @@ -25,6 +25,7 @@ import dcargs import torch from torch import nn from torch.nn import Parameter +from torch.utils.data import Dataset from torch.utils.data.distributed import DistributedSampler from nerfstudio.cameras.rays import RayBundle @@ -101,16 +102,16 @@ class DataManager(nn.Module): Attributes: train_count (int): the step number of our train iteration, needs to be incremented manually eval_count (int): the step number of our eval iteration, needs to be incremented manually - train_input_dataset (InputDataset): the input dataset for the train dataset - eval_input_dataset (InputDataset): the input dataset for the eval dataset + train_dataset (Dataset): the dataset for the train dataset + eval_dataset (Dataset): the dataset for the eval dataset Additional attributes specific to each subclass are defined in the setup_train and setup_eval functions. """ - train_input_dataset: Optional[InputDataset] = None - eval_input_dataset: Optional[InputDataset] = None + train_dataset: Optional[Dataset] = None + eval_dataset: Optional[Dataset] = None train_sampler: Optional[DistributedSampler] = None eval_sampler: Optional[DistributedSampler] = None @@ -126,9 +127,9 @@ class DataManager(nn.Module): super().__init__() self.train_count = 0 self.eval_count = 0 - if self.train_input_dataset and self.train_input_dataset.dataset_inputs: + if self.train_dataset: self.setup_train() - if self.eval_input_dataset and self.eval_input_dataset.dataset_inputs: + if self.eval_dataset: self.setup_eval() def forward(self): @@ -262,6 +263,8 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method """ config: VanillaDataManagerConfig + train_dataset: InputDataset + eval_dataset: InputDataset def __init__( self, @@ -278,23 +281,23 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method self.local_rank = local_rank self.sampler = None - self.train_input_dataset = InputDataset(config.dataparser.setup().get_dataset_inputs(split="train")) - self.eval_input_dataset = InputDataset( + self.train_dataset = InputDataset(config.dataparser.setup().get_dataset_inputs(split="train")) + self.eval_dataset = InputDataset( config.dataparser.setup().get_dataset_inputs(split="val" if not test_mode else "test") ) super().__init__() def setup_train(self): """Sets up the data loaders for training""" - assert self.train_input_dataset is not None + assert self.train_dataset is not None if self.world_size > 1: sampler = DistributedSampler( - self.train_input_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42 + self.train_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42 ) else: sampler = None self.train_image_dataloader = CacheDataloader( - self.train_input_dataset, + self.train_dataset, num_images_to_sample_from=self.config.train_num_images_to_sample_from, device=self.device, num_workers=self.world_size * 4, @@ -303,19 +306,19 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method ) 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_input_dataset.dataset_inputs.cameras.to(self.device)) + self.train_ray_generator = RayGenerator(self.train_dataset.dataset_inputs.cameras.to(self.device)) def setup_eval(self): """Sets up the data loader for evaluation""" - assert self.eval_input_dataset is not None + assert self.eval_dataset is not None if self.world_size > 1: sampler = DistributedSampler( - self.eval_input_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42 + self.eval_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42 ) else: sampler = None self.eval_image_dataloader = CacheDataloader( - self.eval_input_dataset, + self.eval_dataset, num_images_to_sample_from=self.config.eval_num_images_to_sample_from, device=self.device, num_workers=self.world_size * 4, @@ -324,15 +327,15 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method ) self.iter_eval_image_dataloader = iter(self.eval_image_dataloader) self.eval_pixel_sampler = PixelSampler(self.config.eval_num_rays_per_batch) - self.eval_ray_generator = RayGenerator(self.eval_input_dataset.dataset_inputs.cameras.to(self.device)) + self.eval_ray_generator = RayGenerator(self.eval_dataset.dataset_inputs.cameras.to(self.device)) # for loading full images self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader( - input_dataset=self.eval_input_dataset, + input_dataset=self.eval_dataset, device=self.device, num_workers=self.world_size * 4, ) self.eval_dataloader = RandIndicesEvalDataloader( - input_dataset=self.eval_input_dataset, + input_dataset=self.eval_dataset, image_indices=self.config.eval_image_indices, device=self.device, num_workers=self.world_size * 4, diff --git a/nerfstudio/data/utils/__init__.py b/nerfstudio/data/utils/__init__.py new file mode 100644 index 000000000..02655775e --- /dev/null +++ b/nerfstudio/data/utils/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index eca056173..f77087d21 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -144,7 +144,7 @@ class Trainer: def train(self) -> None: """Train the model.""" - assert self.pipeline.datamanager.train_input_dataset is not None, "Missing DatsetInputs" + assert self.pipeline.datamanager.train_dataset is not None, "Missing DatsetInputs" self._init_viewer_scene() with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME): @@ -209,9 +209,9 @@ class Trainer: @check_viewer_enabled def _init_viewer_scene(self) -> None: """Initializes viewer scene with given train dataset""" - assert self.viewer_state and self.pipeline.datamanager.train_input_dataset + assert self.viewer_state and self.pipeline.datamanager.train_dataset self.viewer_state.init_scene( - dataset=self.pipeline.datamanager.train_input_dataset, + dataset=self.pipeline.datamanager.train_dataset, start_train=self.config.viewer.start_train, ) diff --git a/nerfstudio/pipelines/base.py b/nerfstudio/pipelines/base.py index b7195c2eb..784cd2183 100644 --- a/nerfstudio/pipelines/base.py +++ b/nerfstudio/pipelines/base.py @@ -211,10 +211,10 @@ class VanillaPipeline(Pipeline): ) self.datamanager.to(device) # TODO(ethan): get rid of scene_box from the model - assert self.datamanager.train_input_dataset is not None, "Missing input dataset" + assert self.datamanager.train_dataset is not None, "Missing input dataset" self.model = config.model.setup( - scene_box=self.datamanager.train_input_dataset.dataset_inputs.scene_box, - num_train_data=len(self.datamanager.train_input_dataset), + scene_box=self.datamanager.train_dataset.dataset_inputs.scene_box, + num_train_data=len(self.datamanager.train_dataset), ) self.model.to(device)