From d1350e1ecfd7ea599ea9b2b50f8960dccb350302 Mon Sep 17 00:00:00 2001 From: maturk Date: Tue, 27 Jun 2023 14:23:03 +0300 Subject: [PATCH] max num threads as config parameter --- nerfstudio/data/datamanagers/parallel_datamanager.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/nerfstudio/data/datamanagers/parallel_datamanager.py b/nerfstudio/data/datamanagers/parallel_datamanager.py index c2da65fc7..f8060fd79 100644 --- a/nerfstudio/data/datamanagers/parallel_datamanager.py +++ b/nerfstudio/data/datamanagers/parallel_datamanager.py @@ -72,6 +72,9 @@ class ParallelDataManagerConfig(DataManagerConfig): queue_size: int = 2 """Size of shared data queue containing generated ray bundles and batches. If queue_size <= 0, the queue size is infinite.""" + max_thread_workers: Optional[int] = 1 + """Maximum number of threads to use in thread pool executor. If None, automatically + set to harware cpu_count + 4.""" class DataProc(mp.Process): @@ -176,7 +179,7 @@ class ParallelDataManager(DataManager, Generic[TDataset]): proc.start() super().__init__() # Prime the executor with the first batch - self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.config.max_thread_workers) self.batch_fut = self.executor.submit(self.data_q.get) def create_train_dataset(self) -> TDataset: