mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
1b0584f7bd
Summary:
This large diff rewrites a significant portion of Implicitron's config hierarchy. The new hierarchy, and some of the default implementation classes, are as follows:
```
Experiment
data_source: ImplicitronDataSource
dataset_map_provider
data_loader_map_provider
model_factory: ImplicitronModelFactory
model: GenericModel
optimizer_factory: ImplicitronOptimizerFactory
training_loop: ImplicitronTrainingLoop
evaluator: ImplicitronEvaluator
```
1) Experiment (used to be ExperimentConfig) is now a top-level Configurable and contains as members mainly (mostly new) high-level factory Configurables.
2) Experiment's job is to run factories, do some accelerate setup and then pass the results to the main training loop.
3) ImplicitronOptimizerFactory and ImplicitronModelFactory are new high-level factories that create the optimizer, scheduler, model, and stats objects.
4) TrainingLoop is a new configurable that runs the main training loop and the inner train-validate step.
5) Evaluator is a new configurable that TrainingLoop uses to run validation/test steps.
6) GenericModel is not the only model choice anymore. Instead, ImplicitronModelBase (by default instantiated with GenericModel) is a member of Experiment and can be easily replaced by a custom implementation by the user.
All the new Configurables are children of ReplaceableBase, and can be easily replaced with custom implementations.
In addition, I added support for the exponential LR schedule, updated the config files and the test, as well as added a config file that reproduces NERF results and a test to run the repro experiment.
Reviewed By: bottler
Differential Revision: D37723227
fbshipit-source-id: b36bee880d6aa53efdd2abfaae4489d8ab1e8a27
83 lines
3.0 KiB
Python
83 lines
3.0 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the BSD-style license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
from typing import Optional, Tuple
|
|
|
|
from pytorch3d.implicitron.tools.config import (
|
|
registry,
|
|
ReplaceableBase,
|
|
run_auto_creation,
|
|
)
|
|
from pytorch3d.renderer.cameras import CamerasBase
|
|
|
|
from .blender_dataset_map_provider import BlenderDatasetMapProvider # noqa
|
|
from .data_loader_map_provider import DataLoaderMap, DataLoaderMapProviderBase
|
|
from .dataset_map_provider import DatasetMap, DatasetMapProviderBase, Task
|
|
from .json_index_dataset_map_provider import JsonIndexDatasetMapProvider # noqa
|
|
from .json_index_dataset_map_provider_v2 import JsonIndexDatasetMapProviderV2 # noqa
|
|
from .llff_dataset_map_provider import LlffDatasetMapProvider # noqa
|
|
|
|
|
|
class DataSourceBase(ReplaceableBase):
|
|
"""
|
|
Base class for a data source in Implicitron. It encapsulates Dataset
|
|
and DataLoader configuration.
|
|
"""
|
|
|
|
def get_datasets_and_dataloaders(self) -> Tuple[DatasetMap, DataLoaderMap]:
|
|
raise NotImplementedError()
|
|
|
|
@property
|
|
def all_train_cameras(self) -> Optional[CamerasBase]:
|
|
"""
|
|
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.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
def get_task(self) -> Task:
|
|
raise NotImplementedError()
|
|
|
|
|
|
@registry.register
|
|
class ImplicitronDataSource(DataSourceBase): # pyre-ignore[13]
|
|
"""
|
|
Represents the data used in Implicitron. This is the only implementation
|
|
of DataSourceBase provided.
|
|
|
|
Members:
|
|
dataset_map_provider_class_type: identifies type for dataset_map_provider.
|
|
e.g. JsonIndexDatasetMapProvider for Co3D.
|
|
data_loader_map_provider_class_type: identifies type for data_loader_map_provider.
|
|
"""
|
|
|
|
dataset_map_provider: DatasetMapProviderBase
|
|
dataset_map_provider_class_type: str
|
|
data_loader_map_provider: DataLoaderMapProviderBase
|
|
data_loader_map_provider_class_type: str = "SequenceDataLoaderMapProvider"
|
|
|
|
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()
|
|
dataloaders = self.data_loader_map_provider.get_data_loader_map(datasets)
|
|
return datasets, dataloaders
|
|
|
|
def get_task(self) -> Task:
|
|
return self.dataset_map_provider.get_task()
|
|
|
|
@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]
|