diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 891cd272a..6fe32d1bd 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -30,7 +30,7 @@ from typing_extensions import Literal, assert_never from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3 from nerfstudio.configs import base_config as cfg -from nerfstudio.engine.optimizers import AdamOptimizerConfig +from nerfstudio.engine.optimizers import AdamOptimizerConfig, OptimizerConfig from nerfstudio.engine.schedulers import SchedulerConfig from nerfstudio.utils import poses as pose_utils @@ -50,7 +50,7 @@ class CameraOptimizerConfig(cfg.InstantiateConfig): orientation_noise_std: float = 0.0 """Noise to add to initial orientations. Useful for debugging.""" - optimizer: AdamOptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15) + optimizer: OptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15) """ADAM parameters for camera optimization.""" scheduler: SchedulerConfig = SchedulerConfig(max_steps=10000) diff --git a/nerfstudio/data/datamanagers/raycache_datamanager.py b/nerfstudio/data/datamanagers/raycache_datamanager.py new file mode 100644 index 000000000..2c48c25f5 --- /dev/null +++ b/nerfstudio/data/datamanagers/raycache_datamanager.py @@ -0,0 +1,149 @@ +# Copyright 2022 The Nerfstudio 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. + +""" +Data manager for ray level caching +""" + +from __future__ import annotations + +import random +from abc import abstractmethod +from dataclasses import dataclass, field +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 +from torch.utils.data.distributed import DistributedSampler +from typing_extensions import Literal + +import nerfstudio.utils.profiler as profiler +from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig +from nerfstudio.cameras.rays import RayBundle +from nerfstudio.configs.base_config import InstantiateConfig +from nerfstudio.data.datamanagers.base_datamanager import ( + DataManager, + VanillaDataManager, + VanillaDataManagerConfig, +) +from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig +from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig +from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig +from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( + InstantNGPDataParserConfig, +) +from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig +from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig +from nerfstudio.data.dataparsers.phototourism_dataparser import ( + PhototourismDataParserConfig, +) +from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig +from nerfstudio.data.datasets.base_dataset import InputDataset +from nerfstudio.data.pixel_samplers import PixelSampler +from nerfstudio.data.utils.dataloaders import ( + CacheDataloader, + FixedIndicesEvalDataloader, + RandIndicesEvalDataloader, +) +from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate +from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes +from nerfstudio.engine.optimizers import AdamOptimizerConfig +from nerfstudio.model_components.ray_generators import RayGenerator +from nerfstudio.utils.misc import IterableWrapper + +CONSOLE = Console(width=120) + + +@dataclass +class RayCacheDataManagerConfig(VanillaDataManagerConfig): + """A datamanager for optimized caching of ray level data.""" + + _target: Type = field(default_factory=lambda: RayCacheDataManager) + + # train_num_images_to_sample_from: int = 40 + # """Number of images to sample during training iteration.""" + # train_num_times_to_repeat_images: int = 100 + # """When not training on all images, number of iterations before picking new + # images. If -1, never pick new images.""" + # eval_num_images_to_sample_from: int = 40 + # """Number of images to sample during eval iteration.""" + # eval_num_times_to_repeat_images: int = 100 + # """When not training on all images, number of iterations before picking new""" + + # ---------- RayCacheDataManager Specific Args ----------# + num_rays_cached: int = 100000 + """Number of rays to cache / refresh buffer with.""" + train_num_times_to_repeat_ray_cache: int = 50 + """Number of iterations before refreshing the ray cache. + Also ends up being number of steps in between camera updates.""" + camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig(optimizer=AdamOptimizerConfig(weight_decay=0)) + """Camera optimizer config.""" + + +class RayCacheDataManager(VanillaDataManager): # pylint: disable=abstract-method + """ + Wrapper around the vanilla data manager that will cache raw rays + + Nothing is fundamentally flawed or incorrect about the original datamanager... it just could be quicker. + """ + + config: RayCacheDataManagerConfig + + _ray_cache = None + _batch_cache = None + + def _refresh_cache(self): + """Refresh the ray cache and return the next batch.""" + + 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, + collate_fn=self.config.collate_fn, + ) + self.iter_train_image_dataloader = iter(self.train_image_dataloader) + self.train_pixel_sampler = PixelSampler(self.config.num_rays_cached) + self.train_camera_optimizer = self.config.camera_optimizer.setup( + num_cameras=self.train_dataset.cameras.size, device=self.device + ) + self.train_ray_generator = RayGenerator( + self.train_dataset.cameras.to(self.device), + self.train_camera_optimizer, + ) + + def next_train(self, step: int): + """Get the next batch of training data.""" + if self._ray_cache is None or step % self.config.train_num_times_to_repeat_ray_cache == 0: + self._ray_cache, self._batch_cache = super().next_train(step) + + start = random.randint(0, self.config.num_rays_cached - self.config.train_num_rays_per_batch) + rays = self._ray_cache[start : start + self.config.train_num_rays_per_batch] + batch = {k: v[start : start + self.config.train_num_rays_per_batch] for k, v in self._batch_cache.items()} + + if step % self.config.train_num_times_to_repeat_ray_cache == 0: + self._ray_cache = self._ray_cache.detach() + self._batch_cache = {k: v.detach() for k, v in self._batch_cache.items()} + + return rays, batch diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py index 5bae109d0..4595af90c 100644 --- a/nerfstudio/engine/optimizers.py +++ b/nerfstudio/engine/optimizers.py @@ -35,7 +35,6 @@ class OptimizerConfig(base_config.PrintableConfig): _target: Type = torch.optim.Adam lr: float = 0.0005 - eps: float = 1e-08 # TODO: somehow make this more generic. i dont like the idea of overriding the setup function # but also not sure how to go about passing things into predefined torch objects. @@ -46,12 +45,28 @@ class OptimizerConfig(base_config.PrintableConfig): return self._target(params, **kwargs) +@dataclass +class SGDOptimizerConfig(OptimizerConfig): + """Basic optimizer config with RAdam""" + + _target: Type = torch.optim.SGD + lr: float = 0.0005 + weight_decay: float = 0.0 + + def setup(self, params) -> Any: + """Returns the instantiated object using the config.""" + kwargs = vars(self).copy() + kwargs.pop("_target") + return self._target(params, **kwargs) + + @dataclass class AdamOptimizerConfig(OptimizerConfig): """Basic optimizer config with Adam""" _target: Type = torch.optim.Adam weight_decay: float = 0 + eps: float = 1e-08 @dataclass @@ -59,6 +74,7 @@ class RAdamOptimizerConfig(OptimizerConfig): """Basic optimizer config with RAdam""" _target: Type = torch.optim.RAdam + eps: float = 1e-08 def setup_optimizers(config: base_config.Config, param_groups: Dict[str, List[Parameter]]) -> "Optimizers": diff --git a/nerfstudio/utils/tensor_dataclass.py b/nerfstudio/utils/tensor_dataclass.py index 5c8bc26b5..c78e37640 100644 --- a/nerfstudio/utils/tensor_dataclass.py +++ b/nerfstudio/utils/tensor_dataclass.py @@ -147,7 +147,7 @@ class TensorDataclass: return new_dict def __getitem__(self: TensorDataclassT, indices) -> TensorDataclassT: - if isinstance(indices, (torch.Tensor)): + if isinstance(indices, (torch.Tensor, list)): return self._apply_fn_to_fields(lambda x: x[indices]) if isinstance(indices, (int, slice, type(Ellipsis))): indices = (indices,) @@ -256,6 +256,14 @@ class TensorDataclass: """ return self._apply_fn_to_fields(lambda x: x.to(device)) + def detach(self: TensorDataclassT) -> TensorDataclassT: + """Returns a new TensorDataclass with the same data but detached from the computation graph. + + Returns: + A new TensorDataclass with the same data but detached from the computation graph. + """ + return self._apply_fn_to_fields(lambda x: x.detach()) + def _apply_fn_to_fields( self: TensorDataclassT, fn: Callable,