Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

minor patch for exporting pcs and meshes with parallel datamanager #2503

Merged
merged 2 commits into from
Oct 11, 2023
Merged

minor patch for exporting pcs and meshes with parallel datamanager #2503

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9ff7b33
minor patch for exporting pcs and meshes with parallel datamanager
kerrj Oct 10, 2023
3721db1
format
kerrj Oct 10, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions nerfstudio/data/datamanagers/parallel_datamanager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,15 +222,15 @@ def _get_pixel_sampler(self, dataset: TDataset, num_rays_per_batch: int) -> Pixe
def setup_train(self):
"""Sets up parallel python data processes for training."""
assert self.train_dataset is not None
self.train_pix_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch) # type: ignore
self.train_pixel_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch) # type: ignore
self.data_queue = mp.Manager().Queue(maxsize=self.config.queue_size)
self.data_procs = [
DataProcessor(
out_queue=self.data_queue, # type: ignore
config=self.config,
dataparser_outputs=self.train_dataparser_outputs,
dataset=self.train_dataset,
pixel_sampler=self.train_pix_sampler,
pixel_sampler=self.train_pixel_sampler,
)
for i in range(self.config.num_processes)
]
Expand Down Expand Up @@ -319,3 +319,10 @@ def get_param_groups(self) -> Dict[str, List[Parameter]]:
A list of dictionaries containing the data manager's param groups.
"""
return {}

def __del__(self):
"""Clean up the parallel data processes."""
if hasattr(self, "data_procs"):
for proc in self.data_procs:
proc.terminate()
proc.join()
5 changes: 3 additions & 2 deletions nerfstudio/scripts/exporter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@

from nerfstudio.cameras.rays import RayBundle
from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager
from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManager
from nerfstudio.data.scene_box import OrientedBox
from nerfstudio.exporter import texture_utils, tsdf_utils
from nerfstudio.exporter.exporter_utils import (
Expand Down Expand Up @@ -137,7 +138,7 @@ def main(self) -> None:
validate_pipeline(self.normal_method, self.normal_output_name, pipeline)

# Increase the batchsize to speed up the evaluation.
assert isinstance(pipeline.datamanager, VanillaDataManager)
assert isinstance(pipeline.datamanager, (VanillaDataManager, ParallelDataManager))
assert pipeline.datamanager.train_pixel_sampler is not None
pipeline.datamanager.train_pixel_sampler.num_rays_per_batch = self.num_rays_per_batch

Expand Down Expand Up @@ -307,7 +308,7 @@ def main(self) -> None:
validate_pipeline(self.normal_method, self.normal_output_name, pipeline)

# Increase the batchsize to speed up the evaluation.
assert isinstance(pipeline.datamanager, VanillaDataManager)
assert isinstance(pipeline.datamanager, (VanillaDataManager, ParallelDataManager))
assert pipeline.datamanager.train_pixel_sampler is not None
pipeline.datamanager.train_pixel_sampler.num_rays_per_batch = self.num_rays_per_batch

Expand Down