From 3783437d2f1d177d035a7ef43c8138488ac4ccd2 Mon Sep 17 00:00:00 2001 From: Jeremy Reizenstein Date: Thu, 21 Jul 2022 15:04:00 -0700 Subject: [PATCH] lazy all_train_cameras Summary: Avoid calculating all_train_cameras before it is needed, because it is slow in some datasets. Reviewed By: shapovalov Differential Revision: D38037157 fbshipit-source-id: 95461226655cde2626b680661951ab17ebb0ec75 --- projects/implicitron_trainer/experiment.py | 7 +++---- pytorch3d/implicitron/dataset/data_source.py | 15 +++++++++++---- pytorch3d/implicitron/eval_demo.py | 4 +--- tests/implicitron/test_batch_sampler.py | 2 ++ tests/implicitron/test_data_json_index.py | 2 +- tests/implicitron/test_data_llff.py | 2 +- 6 files changed, 19 insertions(+), 13 deletions(-) diff --git a/projects/implicitron_trainer/experiment.py b/projects/implicitron_trainer/experiment.py index 42cc0df1..16bea062 100755 --- a/projects/implicitron_trainer/experiment.py +++ b/projects/implicitron_trainer/experiment.py @@ -391,7 +391,6 @@ def run_training(cfg: DictConfig) -> None: datasource = ImplicitronDataSource(**cfg.data_source_args) datasets, dataloaders = datasource.get_datasets_and_dataloaders() task = datasource.get_task() - all_train_cameras = datasource.get_all_train_cameras() # init the model model, stats, optimizer_state = init_model(cfg=cfg, accelerator=accelerator) @@ -405,7 +404,7 @@ def run_training(cfg: DictConfig) -> None: _eval_and_dump( cfg, task, - all_train_cameras, + datasource.all_train_cameras, datasets, dataloaders, model, @@ -490,7 +489,7 @@ def run_training(cfg: DictConfig) -> None: ): _run_eval( model, - all_train_cameras, + datasource.all_train_cameras, dataloaders.test, task, camera_difficulty_bin_breaks=cfg.camera_difficulty_bin_breaks, @@ -525,7 +524,7 @@ def run_training(cfg: DictConfig) -> None: _eval_and_dump( cfg, task, - all_train_cameras, + datasource.all_train_cameras, datasets, dataloaders, model, diff --git a/pytorch3d/implicitron/dataset/data_source.py b/pytorch3d/implicitron/dataset/data_source.py index 9e12e3c1..a83789cc 100644 --- a/pytorch3d/implicitron/dataset/data_source.py +++ b/pytorch3d/implicitron/dataset/data_source.py @@ -30,9 +30,10 @@ class DataSourceBase(ReplaceableBase): def get_datasets_and_dataloaders(self) -> Tuple[DatasetMap, DataLoaderMap]: raise NotImplementedError() - def get_all_train_cameras(self) -> Optional[CamerasBase]: + @property + def all_train_cameras(self) -> Optional[CamerasBase]: """ - If the data is all for a single scene, returns a list + If the data is all for a single scene, a list of the known training cameras for that scene, which is used for evaluating the viewpoint difficulty of the unseen cameras. @@ -59,6 +60,7 @@ class ImplicitronDataSource(DataSourceBase): # pyre-ignore[13] def __post_init__(self): run_auto_creation(self) + self._all_train_cameras_cache: Optional[Tuple[Optional[CamerasBase]]] = None def get_datasets_and_dataloaders(self) -> Tuple[DatasetMap, DataLoaderMap]: datasets = self.dataset_map_provider.get_dataset_map() @@ -68,5 +70,10 @@ class ImplicitronDataSource(DataSourceBase): # pyre-ignore[13] def get_task(self) -> Task: return self.dataset_map_provider.get_task() - def get_all_train_cameras(self) -> Optional[CamerasBase]: - return self.dataset_map_provider.get_all_train_cameras() + @property + def all_train_cameras(self) -> Optional[CamerasBase]: + if self._all_train_cameras_cache is None: # pyre-ignore[16] + all_train_cameras = self.dataset_map_provider.get_all_train_cameras() + self._all_train_cameras_cache = (all_train_cameras,) + + return self._all_train_cameras_cache[0] diff --git a/pytorch3d/implicitron/eval_demo.py b/pytorch3d/implicitron/eval_demo.py index c8daee5c..b67e1b79 100644 --- a/pytorch3d/implicitron/eval_demo.py +++ b/pytorch3d/implicitron/eval_demo.py @@ -118,8 +118,6 @@ def evaluate_dbir_for_category( if test_dataset is None or test_dataloader is None: raise ValueError("must have a test dataset.") - all_train_cameras = data_source.get_all_train_cameras() - image_size = cast(JsonIndexDataset, test_dataset).image_width if image_size is None: @@ -149,7 +147,7 @@ def evaluate_dbir_for_category( preds["implicitron_render"], bg_color=bg_color, lpips_model=lpips_model, - source_cameras=all_train_cameras, + source_cameras=data_source.all_train_cameras, ) ) diff --git a/tests/implicitron/test_batch_sampler.py b/tests/implicitron/test_batch_sampler.py index 2b64804c..50d0b8fe 100644 --- a/tests/implicitron/test_batch_sampler.py +++ b/tests/implicitron/test_batch_sampler.py @@ -10,6 +10,7 @@ from collections import defaultdict from dataclasses import dataclass from itertools import product +import numpy as np from pytorch3d.implicitron.dataset.data_loader_map_provider import ( DoublePoolBatchSampler, ) @@ -53,6 +54,7 @@ class MockDataset(DatasetBase): class TestSceneBatchSampler(unittest.TestCase): def setUp(self): + np.random.seed(42) self.dataset_overfit = MockDataset(1) def test_overfit(self): diff --git a/tests/implicitron/test_data_json_index.py b/tests/implicitron/test_data_json_index.py index a07cd13f..0cd77c82 100644 --- a/tests/implicitron/test_data_json_index.py +++ b/tests/implicitron/test_data_json_index.py @@ -31,7 +31,7 @@ class TestDataJsonIndex(TestCaseMixin, unittest.TestCase): data_source = ImplicitronDataSource(**args) - cameras = data_source.get_all_train_cameras() + cameras = data_source.all_train_cameras self.assertIsInstance(cameras, PerspectiveCameras) self.assertEqual(len(cameras), 81) diff --git a/tests/implicitron/test_data_llff.py b/tests/implicitron/test_data_llff.py index 9f73478b..032b2f03 100644 --- a/tests/implicitron/test_data_llff.py +++ b/tests/implicitron/test_data_llff.py @@ -152,6 +152,6 @@ class TestDataLlff(TestCaseMixin, unittest.TestCase): self.assertEqual(i.frame_type, ["unseen"]) self.assertEqual(i.image_rgb.shape, (1, 3, 800, 800)) - cameras = data_source.get_all_train_cameras() + cameras = data_source.all_train_cameras self.assertIsInstance(cameras, PerspectiveCameras) self.assertEqual(len(cameras), 100)