From 1ce3dd2f75a71848f3296459c40a138f4c2d0670 Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Sun, 6 Nov 2022 21:40:39 -0800 Subject: [PATCH] better memory --- .../data/datamanagers/base_datamanager.py | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 9d82eb2ab..b721d7b77 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -265,7 +265,9 @@ class VanillaDataManagerConfig(InstantiateConfig): Record3D.""" import torch.multiprocessing as mp -MP_CON = mp.get_context('spawn') +import queue +import time +MP_CON = mp.get_context('forkserver') class Runner(MP_CON.Process): def __init__(self, bundle_queue, train_dataset, device, config): super().__init__() @@ -294,8 +296,14 @@ class Runner(MP_CON.Process): image_batch = next(self.iter_train_image_dataloader) batch = self.train_pixel_sampler.sample(image_batch) ray_indices = batch["indices"] - ray_bundle = self.train_ray_generator(ray_indices) - self.bundle_queue.put((ray_bundle,batch),block=True) + ray_bundle = self.train_ray_generator(ray_indices).to(self.device) + while True: + try: + self.bundle_queue.put((ray_bundle,batch),block=False) + break + except queue.Full: + time.sleep(.02) + class VanillaDataManager(DataManager): # pylint: disable=abstract-method """Basic stored data manager implementation. @@ -332,7 +340,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method self.train_dataset = self.create_train_dataset() self.eval_dataset = self.create_eval_dataset() - self.bundle_queue = MP_CON.Queue(5) + self.bundle_queue = MP_CON.Queue(4)#including Manager() is important for speed runners = [Runner(self.bundle_queue, self.train_dataset, self.device,self.config) for i in range(3)] DataManager.__init__(self) @@ -387,12 +395,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method def next_train(self, step: int) -> Tuple[RayBundle, Dict]: """Returns the next batch of data from the train dataloader.""" self.train_count += 1 - # image_batch = next(self.iter_train_image_dataloader) - # batch = self.train_pixel_sampler.sample(image_batch) - # ray_indices = batch["indices"] - # ray_bundle = self.train_ray_generator(ray_indices) ray_bundle,batch = self.bundle_queue.get(block=True) - ray_bundle = ray_bundle.to(self.device) ray_indices = ray_bundle.camera_indices # Multiply in the camera optimization deltas to the rays