From 09e9a71d3da7c15396f8eafd6b6becf40d6d4941 Mon Sep 17 00:00:00 2001 From: Justin Date: Wed, 21 Jun 2023 20:46:20 +0000 Subject: [PATCH] minor speedup --- nerfstudio/configs/method_configs.py | 2 +- .../data/datamanagers/parallel_datamanager.py | 47 ++++++++++--------- 2 files changed, 26 insertions(+), 23 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 3d64953a9..8c10ad579 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -88,7 +88,7 @@ method_configs["nerfacto"] = TrainerConfig( mixed_precision=True, pipeline=VanillaPipelineConfig( datamanager=ParallelDataManagerConfig( - # datamanager=VanillaDataManagerConfig( + # datamanager=VanillaDataManagerConfig( dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, diff --git a/nerfstudio/data/datamanagers/parallel_datamanager.py b/nerfstudio/data/datamanagers/parallel_datamanager.py index 03ab8f5c5..60affa17b 100644 --- a/nerfstudio/data/datamanagers/parallel_datamanager.py +++ b/nerfstudio/data/datamanagers/parallel_datamanager.py @@ -2,16 +2,15 @@ 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) +from typing import Any, Callable, Dict, Generic, List, Literal, Optional, Tuple, Type, Union, cast import torch +import torch.multiprocessing as mp from rich.progress import track from torch import nn from torch.nn import Parameter @@ -22,20 +21,14 @@ 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.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.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.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 @@ -76,6 +69,7 @@ class ParallelDataManagerConfig(DataManagerConfig): n_procs: int = 1 """Number of processes to use for data loading.""" + class DataProc(mp.Process): def __init__( self, @@ -103,14 +97,19 @@ class DataProc(mp.Process): 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: RayBundle = self.train_ray_generator(ray_indices) ray_bundle = ray_bundle.pin_memory() - self.out_queue.put((ray_bundle,batch), block=True) + while True: + try: + self.out_queue.put_nowait((ray_bundle, batch)) + break + except: + time.sleep(0.005) def cache_images(self): # caches all the input images in a NxHxWx3 tensor @@ -157,8 +156,11 @@ class ParallelDataManager(DataManager, Generic[TDataset]): 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) + # Spawn is critical for not freezing the program (PyTorch compatability issue) + mp.set_start_method("spawn") + # Manager().queue() can be faster + # https://stackoverflow.com/questions/43439194/python-multiprocessing-queue-vs-multiprocessing-manager-queue/45236748#45236748 + self.data_q = mp.Manager().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) @@ -190,16 +192,17 @@ class ParallelDataManager(DataManager, Generic[TDataset]): 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 - + return bundle, batch + def get_train_rays_per_batch(self) -> int: return self.config.train_num_rays_per_batch