From e26bd4d52530d7b67c401ac94025ba065bef7126 Mon Sep 17 00:00:00 2001 From: Justin Date: Wed, 21 Jun 2023 18:50:22 +0000 Subject: [PATCH] initial impl of parallel data manager --- nerfstudio/configs/method_configs.py | 12 +- .../data/datamanagers/base_datamanager.py | 26 +-- .../data/datamanagers/parallel_datamanager.py | 217 ++++++++++++++++++ nerfstudio/data/pixel_samplers.py | 10 +- nerfstudio/fields/nerfacto_field.py | 31 ++- nerfstudio/models/nerfacto.py | 10 +- nerfstudio/utils/tensor_dataclass.py | 8 + .../viewer_beta/render_state_machine.py | 4 +- nerfstudio/viewer_beta/viewer.py | 2 +- 9 files changed, 280 insertions(+), 40 deletions(-) create mode 100644 nerfstudio/data/datamanagers/parallel_datamanager.py diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 405bac6a2..3d64953a9 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -27,6 +27,7 @@ from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.configs.base_config import ViewerConfig from nerfstudio.configs.external_methods import get_external_methods from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager, VanillaDataManagerConfig +from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManagerConfig from nerfstudio.data.datamanagers.random_cameras_datamanager import RandomCamerasDataManagerConfig from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig @@ -86,7 +87,8 @@ method_configs["nerfacto"] = TrainerConfig( max_num_iterations=30000, mixed_precision=True, pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig( + datamanager=ParallelDataManagerConfig( + # datamanager=VanillaDataManagerConfig( dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, @@ -107,7 +109,7 @@ method_configs["nerfacto"] = TrainerConfig( }, "camera_opt": { "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15), - "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=2000), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=5000), }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), @@ -199,12 +201,12 @@ method_configs["nerfacto-huge"] = TrainerConfig( "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=50000), }, "camera_opt": { - "optimizer": RAdamOptimizerConfig(lr=6e-4, eps=1e-8), - "scheduler": ExponentialDecaySchedulerConfig(lr_final=6e-5, max_steps=50000), + "optimizer": RAdamOptimizerConfig(lr=1e-3, eps=1e-8), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=10000), }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), - vis="viewer", + vis="viewer_beta", ) # method_configs["depth-nerfacto"] = TrainerConfig( diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index c957af0ca..580f8f3bf 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -23,19 +23,7 @@ from abc import abstractmethod from collections import defaultdict from dataclasses import dataclass, field from pathlib import Path -from typing import ( - Any, - Callable, - Dict, - Generic, - List, - Literal, - Optional, - Tuple, - Type, - Union, - cast, -) +from typing import Any, Callable, Dict, Generic, List, Literal, Optional, Tuple, Type, Union, cast import torch from torch import nn @@ -50,16 +38,8 @@ from nerfstudio.configs.dataparser_configs import AnnotatedDataParserUnion from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset -from nerfstudio.data.pixel_samplers import ( - EquirectangularPixelSampler, - PatchPixelSampler, - PixelSampler, -) -from nerfstudio.data.utils.dataloaders import ( - CacheDataloader, - FixedIndicesEvalDataloader, - RandIndicesEvalDataloader, -) +from nerfstudio.data.pixel_samplers import EquirectangularPixelSampler, PatchPixelSampler, 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.model_components.ray_generators import RayGenerator diff --git a/nerfstudio/data/datamanagers/parallel_datamanager.py b/nerfstudio/data/datamanagers/parallel_datamanager.py new file mode 100644 index 000000000..03ab8f5c5 --- /dev/null +++ b/nerfstudio/data/datamanagers/parallel_datamanager.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +import concurrent.futures +import functools +import torch.multiprocessing as mp +import time +from abc import abstractmethod +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import (Any, Callable, Dict, Generic, List, Literal, Optional, + Tuple, Type, Union, cast) + +import torch +from rich.progress import track +from torch import nn +from torch.nn import Parameter +from torch.utils.data.distributed import DistributedSampler +from typing_extensions import TypeVar + +from nerfstudio.cameras.cameras import CameraType +from nerfstudio.cameras.rays import RayBundle +from nerfstudio.configs.base_config import InstantiateConfig +from nerfstudio.configs.dataparser_configs import AnnotatedDataParserUnion +from nerfstudio.data.datamanagers.base_datamanager import ( + DataManager, DataManagerConfig, TDataset, variable_res_collate) +from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs +from nerfstudio.data.dataparsers.blender_dataparser import \ + BlenderDataParserConfig +from nerfstudio.data.datasets.base_dataset import InputDataset +from nerfstudio.data.pixel_samplers import (EquirectangularPixelSampler, + PatchPixelSampler, 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.model_components.ray_generators import RayGenerator +from nerfstudio.utils.misc import IterableWrapper +from nerfstudio.utils.rich_utils import CONSOLE + + +@dataclass +class ParallelDataManagerConfig(DataManagerConfig): + """A basic data manager""" + + _target: Type = field(default_factory=lambda: ParallelDataManager) + """Target class to instantiate.""" + dataparser: AnnotatedDataParserUnion = BlenderDataParserConfig() + """Specifies the dataparser used to unpack the data.""" + train_num_rays_per_batch: int = 1024 + """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.""" + collate_fn: Callable[[Any], Any] = cast(Any, staticmethod(nerfstudio_collate)) + """Specifies the collate function to use for the train and eval dataloaders.""" + camera_res_scale_factor: float = 1.0 + """The scale factor for scaling spatial data such as images, mask, semantics + along with relevant information about camera intrinsics + """ + patch_size: int = 1 + """Size of patch to sample from. If >1, patch-based sampling will be used.""" + n_procs: int = 1 + """Number of processes to use for data loading.""" + +class DataProc(mp.Process): + def __init__( + self, + out_queue: mp.Queue, + config: ParallelDataManagerConfig, + dataparser_outputs: DataparserOutputs, + train_dataset: TDataset, + pix_sampler: PixelSampler, + ): + super().__init__() + self.daemon = True + self.out_queue = out_queue + self.config = config + self.dataparser_outputs = dataparser_outputs + self.train_dataset = train_dataset + self.exclude_batch_keys_from_device = self.train_dataset.exclude_batch_keys_from_device + if self.dataparser_outputs is not None: + cameras = self.dataparser_outputs.cameras + if len(cameras) > 1: + for i in range(1, len(cameras)): + if cameras[0].width != cameras[i].width or cameras[0].height != cameras[i].height: + CONSOLE.print("Variable resolution, using variable_res_collate") + self.config.collate_fn = variable_res_collate + break + self.train_pixel_sampler = pix_sampler + self.train_ray_generator = RayGenerator(self.train_dataset.cameras) + self.cache_images() + + def run(self): + while True: + batch = self.train_pixel_sampler.sample(self.img_data) + ray_indices = batch["indices"] + ray_bundle:RayBundle = self.train_ray_generator(ray_indices) + ray_bundle = ray_bundle.pin_memory() + self.out_queue.put((ray_bundle,batch), block=True) + + def cache_images(self): + # caches all the input images in a NxHxWx3 tensor + indices = range(len(self.train_dataset)) + batch_list = [] + results = [] + with concurrent.futures.ThreadPoolExecutor() as executor: + for idx in indices: + res = executor.submit(self.train_dataset.__getitem__, idx) + results.append(res) + for res in track(results, description="Loading data batch", transient=True): + batch_list.append(res.result()) + self.img_data = self.config.collate_fn(batch_list) + + +class ParallelDataManager(DataManager, Generic[TDataset]): + def __init__( + self, + config: ParallelDataManagerConfig, + device: Union[torch.device, str] = "cpu", + test_mode: Literal["test", "val", "inference"] = "val", + world_size: int = 1, + local_rank: int = 0, + **kwargs, + ): + self.dataset_type: Type[TDataset] = kwargs.get("_dataset_type", getattr(TDataset, "__default__")) + self.config = config + self.device = device + self.world_size = world_size + self.local_rank = local_rank + # self.sampler = None + self.test_mode = test_mode + self.test_split = "test" if test_mode in ["test", "inference"] else "val" + self.dataparser_config = self.config.dataparser + if self.config.data is not None: + self.config.dataparser.data = Path(self.config.data) + else: + self.config.data = self.config.dataparser.data + self.dataparser = self.dataparser_config.setup() + if test_mode == "inference": + self.dataparser.downscale_factor = 1 # Avoid opening images + self.includes_time = self.dataparser.includes_time + self.train_dataparser_outputs: DataparserOutputs = self.dataparser.get_dataparser_outputs(split="train") + self.train_dataset = self.create_train_dataset() + # self.eval_dataset = self.create_eval_dataset() + self.pix_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch) + mp.set_start_method('spawn') + self.data_q = mp.Queue(maxsize=4) + self.data_procs = [ + DataProc(self.data_q, self.config, self.train_dataparser_outputs, self.train_dataset, self.pix_sampler) + for i in range(self.config.n_procs) + ] + for proc in self.data_procs: + proc.start() + super().__init__() + # Prime the executor with the first batch + self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + self.batch_fut = self.executor.submit(self.data_q.get) + + def create_train_dataset(self) -> TDataset: + """Sets up the data loaders for training""" + return self.dataset_type( + dataparser_outputs=self.train_dataparser_outputs, + scale_factor=self.config.camera_res_scale_factor, + ) + + def _get_pixel_sampler(self, dataset: TDataset, *args: Any, **kwargs: Any) -> PixelSampler: + """Infer pixel sampler to use.""" + if self.config.patch_size > 1: + return PatchPixelSampler(*args, **kwargs, patch_size=self.config.patch_size) + + # If all images are equirectangular, use equirectangular pixel sampler + is_equirectangular = dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value + if is_equirectangular.all(): + return EquirectangularPixelSampler(*args, **kwargs) + # Otherwise, use the default pixel sampler + if is_equirectangular.any(): + CONSOLE.print("[bold yellow]Warning: Some cameras are equirectangular, but using default pixel sampler.") + return PixelSampler(*args, **kwargs) + + def next_train(self, step: int) -> Tuple[RayBundle, Dict]: + self.train_count += 1 + # Fetch the next batch in an executor to parallelize the queue get() operation + # with the train step + bundle, batch = self.batch_fut.result() + self.batch_fut = self.executor.submit(self.data_q.get) + bundle = bundle.to(self.device) + return bundle,batch + + def get_train_rays_per_batch(self) -> int: + return self.config.train_num_rays_per_batch + + def get_eval_rays_per_batch(self) -> int: + return self.config.eval_num_rays_per_batch + + def get_datapath(self) -> Path: + return self.config.dataparser.data + + def get_param_groups(self) -> Dict[str, List[Parameter]]: + """Get the param groups for the data manager. + Returns: + A list of dictionaries containing the data manager's param groups. + """ + return {} diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index a145dce46..a51363a54 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -67,10 +67,10 @@ class PixelSampler: chosen_indices = random.sample(range(len(nonzero_indices)), k=batch_size) indices = nonzero_indices[chosen_indices] else: - indices = torch.floor( - torch.rand((batch_size, 3), device=device) - * torch.tensor([num_images, image_height, image_width], device=device) - ).long() + randids = torch.rand((batch_size, 3), device=device)* torch.tensor([num_images, image_height, image_width], device=device) + # indices = randids.floor() + # print("done floor") + indices = randids.long() return indices @@ -101,13 +101,11 @@ class PixelSampler: collated_batch = { key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None } - assert collated_batch["image"].shape[0] == num_rays_per_batch # Needed to correct the random indices to their actual camera idx locations. indices[:, 0] = batch["image_idx"][c] collated_batch["indices"] = indices # with the abs camera indices - if keep_full_image: collated_batch["full_image"] = batch["image"] diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 1c35b7121..dd3f44df0 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -89,6 +89,7 @@ class NerfactoField(Field): use_transient_embedding: bool = False, use_semantics: bool = False, num_semantic_classes: int = 100, + freq_warmup: int = 0, pass_semantic_gradients: bool = False, use_pred_normals: bool = False, use_average_appearance_embedding: bool = False, @@ -114,6 +115,8 @@ class NerfactoField(Field): self.use_pred_normals = use_pred_normals self.pass_semantic_gradients = pass_semantic_gradients self.base_res = base_res + self.freq_warmup = freq_warmup + self.step = 0 self.direction_encoding = SHEncoding( levels=4, @@ -198,6 +201,29 @@ class NerfactoField(Field): implementation=implementation, ) + def step_cb(self, step): + self.step = step + + def get_freq_mask(self, B, N, D, device): + import numpy as np + + ts = np.linspace(self.step * N, self.step, N) + mask_vals = np.interp(ts, [0, self.freq_warmup], [0, 1]) + m = torch.from_numpy(mask_vals)[None, ..., None].to(device) + return m.repeat((B, 1, D)) + + def mask_freqs(self, hashgrid_outputs): + """ + hashgrid_outputs is a B x N*D array of the hashgrid values + """ + # N is number of levels (usually 16) + # D is feature dim (usually 2) + N = self.mlp_base_grid.num_levels + D = self.mlp_base_grid.features_per_level + hashgrid_outputs = hashgrid_outputs.view(-1, N, D) + masked = self.get_freq_mask(hashgrid_outputs.shape[0], N, D, hashgrid_outputs.device) * hashgrid_outputs + return masked.view(masked.shape[0], -1) + def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: """Computes and returns the densities.""" if self.spatial_distortion is not None: @@ -213,7 +239,10 @@ class NerfactoField(Field): if not self._sample_locations.requires_grad: self._sample_locations.requires_grad = True positions_flat = positions.view(-1, 3) - h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1) + hashgrid_vecs = self.mlp_base_grid(positions_flat) + hashgrid_vecs = self.mask_freqs(hashgrid_vecs) + h = self.mlp_base_mlp(hashgrid_vecs).view(*ray_samples.frustums.shape, -1) + # h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1) density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) self._density_before_activation = density_before_activation diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 0d271a77d..f371c9df8 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -18,6 +18,7 @@ NeRF implementation that combines many recent advancements. from __future__ import annotations +from contextlib import nullcontext from dataclasses import dataclass, field from typing import Dict, List, Literal, Tuple, Type @@ -283,13 +284,19 @@ class NerfactoModel(Model): func=self.proposal_sampler.step_cb, ) ) + callbacks.append( + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + update_every_num_iters=1, + func=self.field.step_cb, + ) + ) return callbacks def get_outputs(self, ray_bundle: RayBundle): # apply the camera optimizer pose tweaks if self.training: self.camera_optimizer.apply_to_raybundle(ray_bundle) - ray_samples: RaySamples ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns) field_outputs = self.field.forward(ray_samples, compute_normals=self.config.predict_normals) @@ -333,7 +340,6 @@ class NerfactoModel(Model): for i in range(self.config.num_proposal_iterations): outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i]) - return outputs def get_metrics_dict(self, outputs, batch): diff --git a/nerfstudio/utils/tensor_dataclass.py b/nerfstudio/utils/tensor_dataclass.py index 0ebe7b5ce..3237cb38a 100644 --- a/nerfstudio/utils/tensor_dataclass.py +++ b/nerfstudio/utils/tensor_dataclass.py @@ -260,6 +260,14 @@ class TensorDataclass: A new TensorDataclass with the same data but on the specified device. """ return self._apply_fn_to_fields(lambda x: x.to(device)) + + def pin_memory(self: TensorDataclassT) -> TensorDataclassT: + """Pins the tensor dataclass memory + + Returns: + TensorDataclass: A new TensorDataclass with the same data but pinned. + """ + return self._apply_fn_to_fields(lambda x: x.pin_memory()) def _apply_fn_to_fields( self: TensorDataclassT, diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index bfbf9570a..9971338c8 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/nerfstudio/viewer_beta/render_state_machine.py @@ -181,8 +181,8 @@ class RenderStateMachine(threading.Thread): # if we got interrupted, don't send the output to the viewer continue self._send_output_to_viewer(outputs) - if self.state == "high": - self.viewer.update_camera_poses() + # if self.state == "high": + # self.viewer.update_camera_poses() # if we rendered a static low res, we need to self-trigger a static high-res if self.next_action is None: # if there hasn't been an action, wait for 1/target_fps seconds in case we get another move command diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index 8621d3659..1e51fc8e9 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/nerfstudio/viewer_beta/viewer.py @@ -212,7 +212,7 @@ class Viewer: R = vtf.SO3.from_matrix(c2w[:3, :3]) R = R @ vtf.SO3.from_x_radians(np.pi) camera_handle = self.viser_server.add_camera_frustum( - name=f"/camera_{idx:05d}", + name=f"/Train Cameras/camera_{idx:05d}", fov=2 * np.arctan(camera.cx / camera.fx[0]), scale=1, aspect=float(camera.cx[0] / camera.cy[0]),