From 0a8ebf93f14fe578ec0621e8514710eafc837a20 Mon Sep 17 00:00:00 2001 From: dozeri83 <109915435+dozeri83@users.noreply.github.com> Date: Tue, 22 Nov 2022 21:18:04 +0200 Subject: [PATCH] implemented threadpool mechanics when loading images (#994) Co-authored-by: dror Co-authored-by: Matthew Tancik --- nerfstudio/data/utils/dataloaders.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index 51d7535b0..0c767121d 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/nerfstudio/data/utils/dataloaders.py @@ -16,6 +16,9 @@ Code for sampling images from a dataset of images. """ +# for multithreading +import concurrent.futures +import multiprocessing import random from abc import abstractmethod from typing import Dict, Optional, Tuple, Union @@ -57,6 +60,7 @@ class CacheDataloader(DataLoader): self.cache_all_images = (num_images_to_sample_from == -1) or (num_images_to_sample_from >= len(self.dataset)) self.num_images_to_sample_from = len(self.dataset) if self.cache_all_images else num_images_to_sample_from self.device = device + self.num_workers = kwargs.get("num_workers", 0) self.num_repeated = self.num_times_to_repeat_images # starting value self.first_time = True @@ -85,10 +89,23 @@ class CacheDataloader(DataLoader): def _get_batch_list(self): """Returns a list of batches from the dataset attribute.""" + indices = random.sample(range(len(self.dataset)), k=self.num_images_to_sample_from) batch_list = [] - for idx in track(indices, description="Loading data batch"): - batch_list.append(self.dataset.__getitem__(idx)) + results = [] + + num_threads = int(self.num_workers) * 4 + num_threads = min(num_threads, multiprocessing.cpu_count() - 1) + num_threads = max(num_threads, 1) + + with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor: + for idx in indices: + res = executor.submit(self.dataset.__getitem__, idx) + results.append(res) + + for res in track(results, description="Loading data batch"): + batch_list.append(res.result()) + return batch_list def _get_collated_batch(self):