tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Abstractionv2 (#500)

Browse Source 
* Pushing abstraction changes

* Sorting imports
This commit is contained in:
Jake Austin 2022-09-20 23:56:36 -07:00 committed by GitHub
parent 70f1b127f3
commit 7cf97348ee
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 130 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
nerfactory/configs/base.py
View File

@ -34,7 +34,7 @@ from nerfactory.models.base import Model
from nerfactory.models.nerfw import NerfWModel
from nerfactory.models.tensorf import TensoRFModel
from nerfactory.optimizers.schedulers import ExponentialDecaySchedule
from nerfactory.pipelines.base import Pipeline
from nerfactory.pipelines.base import VanillaPipeline
from nerfactory.utils import writer
@ -296,10 +296,10 @@ class TensoRFModelConfig(VanillaModelConfig):
# Pipeline related configs
@dataclass
class PipelineConfig(InstantiateConfig):
class VanillaPipelineConfig(InstantiateConfig):
"""Configuration for pipeline instantiation"""
_target: Type = Pipeline
_target: Type = VanillaPipeline
"""target class to instantiate"""
datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig()
"""specifies the datamanager config"""
@ -390,7 +390,7 @@ class Config(PrintableConfig):
logging: LoggingConfig = LoggingConfig()
viewer: ViewerConfig = ViewerConfig()
trainer: TrainerConfig = TrainerConfig()
pipeline: PipelineConfig = PipelineConfig()
pipeline: VanillaPipelineConfig = VanillaPipelineConfig()
optimizers: Dict[str, Any] = to_immutable_dict(
{
"fields": {

18
nerfactory/configs/base_configs.py
View File

@ -29,13 +29,13 @@ from nerfactory.configs.base import (
LoggingConfig,
MipNerf360DataParserConfig,
NerfWModelConfig,
PipelineConfig,
RAdamOptimizerConfig,
SchedulerConfig,
TensoRFModelConfig,
TrainerConfig,
VanillaDataManagerConfig,
VanillaModelConfig,
VanillaPipelineConfig,
ViewerConfig,
)
from nerfactory.datamanagers.dataparsers.friends_parser import FriendsDataParserConfig
@ -50,7 +50,7 @@ base_configs: Dict[str, Config] = {}
base_configs["compound"] = Config(
method_name="compound",
trainer=TrainerConfig(mixed_precision=True),
pipeline=PipelineConfig(
pipeline=VanillaPipelineConfig(
datamanager=VanillaDataManagerConfig(dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192),
model=CompoundModelConfig(eval_num_rays_per_chunk=8192),
),
@ -67,7 +67,7 @@ base_configs["compound"] = Config(
base_configs["instant-ngp"] = Config(
method_name="instant-ngp",
trainer=TrainerConfig(steps_per_eval_batch=500, steps_per_save=2000, mixed_precision=True),
pipeline=PipelineConfig(
pipeline=VanillaPipelineConfig(
datamanager=VanillaDataManagerConfig(dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192),
model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192),
),
@ -84,7 +84,7 @@ base_configs["instant-ngp"] = Config(
base_configs["mipnerf-360"] = Config(
method_name="mipnerf-360",
trainer=TrainerConfig(steps_per_eval_batch=200),
pipeline=PipelineConfig(
pipeline=VanillaPipelineConfig(
datamanager=VanillaDataManagerConfig(dataparser=MipNerf360DataParserConfig(), train_num_rays_per_batch=8192),
model=VanillaModelConfig(
_target=MipNerf360Model,
@ -105,7 +105,7 @@ base_configs["mipnerf-360"] = Config(
base_configs["mipnerf"] = Config(
method_name="mipnerf",
pipeline=PipelineConfig(
pipeline=VanillaPipelineConfig(
datamanager=VanillaDataManagerConfig(dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192),
model=VanillaModelConfig(
_target=MipNerfModel,
@ -125,7 +125,7 @@ base_configs["mipnerf"] = Config(
base_configs["nerfw"] = Config(
method_name="nerfw",
pipeline=PipelineConfig(
pipeline=VanillaPipelineConfig(
datamanager=VanillaDataManagerConfig(
dataparser=FriendsDataParserConfig(),
),
@ -142,7 +142,7 @@ base_configs["nerfw"] = Config(
base_configs["semantic-nerf"] = Config(
method_name="semantic-nerf",
pipeline=PipelineConfig(
pipeline=VanillaPipelineConfig(
datamanager=VanillaDataManagerConfig(
dataparser=FriendsDataParserConfig(),
),
@ -163,7 +163,7 @@ base_configs["semantic-nerf"] = Config(
base_configs["vanilla-nerf"] = Config(
method_name="vanilla-nerf",
pipeline=PipelineConfig(
pipeline=VanillaPipelineConfig(
datamanager=VanillaDataManagerConfig(
dataparser=BlenderDataParserConfig(),
),
@ -180,7 +180,7 @@ base_configs["vanilla-nerf"] = Config(
base_configs["tensorf"] = Config(
method_name="tensorf",
trainer=TrainerConfig(mixed_precision=True),
pipeline=PipelineConfig(
pipeline=VanillaPipelineConfig(
datamanager=VanillaDataManagerConfig(
dataparser=BlenderDataParserConfig(),
),

6
nerfactory/datamanagers/base.py
View File

@ -28,7 +28,7 @@ from torch.utils.data.distributed import DistributedSampler
from nerfactory.cameras.rays import RayBundle
from nerfactory.configs import base as cfg
from nerfactory.datamanagers.dataloaders import (
CacheImageDataloader,
CacheDataloader,
FixedIndicesEvalDataloader,
RandIndicesEvalDataloader,
)
@ -246,7 +246,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method
)
else:
sampler = None
self.train_image_dataloader = CacheImageDataloader(
self.train_image_dataloader = CacheDataloader(
self.train_input_dataset,
num_images_to_sample_from=self.config.train_num_images_to_sample_from,
device=self.device,
@ -267,7 +267,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method
)
else:
sampler = None
self.eval_image_dataloader = CacheImageDataloader(
self.eval_image_dataloader = CacheDataloader(
self.eval_input_dataset,
num_images_to_sample_from=self.config.eval_num_images_to_sample_from,
device=self.device,

8
nerfactory/datamanagers/dataloaders.py
View File

@ -21,7 +21,7 @@ from abc import abstractmethod
from typing import Dict, Optional, Tuple, Union
import torch
from torch.utils.data import default_collate
from torch.utils.data import Dataset, default_collate
from torch.utils.data.dataloader import DataLoader
from nerfactory.cameras.cameras import Cameras
@ -30,8 +30,8 @@ from nerfactory.datamanagers.datasets import InputDataset
from nerfactory.utils.misc import get_dict_to_torch
class CacheImageDataloader(DataLoader):
"""Collated image dataset that implements caching of images.
class CacheDataloader(DataLoader):
"""Collated image dataset that implements caching of default-pytorch-collatable data.
Creates batches of the InputDataset return type.
Args:
@ -43,7 +43,7 @@ class CacheImageDataloader(DataLoader):
def __init__(
self,
dataset: InputDataset,
dataset: Dataset,
num_images_to_sample_from: int = -1,
num_times_to_repeat_images: int = 0,
device: Union[torch.device, str] = "cpu",

6
nerfactory/engine/trainer.py
View File

@ -31,7 +31,7 @@ from torch.cuda.amp.grad_scaler import GradScaler
from nerfactory.configs import base as cfg
from nerfactory.optimizers.optimizers import Optimizers, setup_optimizers
from nerfactory.pipelines.base import Pipeline
from nerfactory.pipelines.base import VanillaPipeline
from nerfactory.utils import profiler, writer
from nerfactory.utils.callbacks import (
TrainingCallback,
@ -79,7 +79,7 @@ class Trainer:
callbacks: The callbacks object.
"""
pipeline: Pipeline
pipeline: VanillaPipeline
optimizers: Optimizers
callbacks: List[TrainingCallback]
@ -119,7 +119,7 @@ class Trainer:
Args:
test_mode: Whether to setup for testing. Defaults to False.
"""
self.pipeline: Pipeline = self.config.pipeline.setup(
self.pipeline: VanillaPipeline = self.config.pipeline.setup(
device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank
)
self.optimizers = setup_optimizers(self.config.optimizers, self.pipeline.get_param_groups())

123
nerfactory/pipelines/base.py
View File

@ -36,7 +36,7 @@ from torch.nn import Parameter
from torch.nn.parallel import DistributedDataParallel as DDP
from nerfactory.configs import base as cfg
from nerfactory.datamanagers.base import DataManager
from nerfactory.datamanagers.base import DataManager, VanillaDataManager
from nerfactory.models.base import Model
from nerfactory.utils import profiler
from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes
@ -69,19 +69,6 @@ class Pipeline(nn.Module):
simplifying the model classes, which each may have different forward() methods
and so on.
TODO: For viewer functionality to be added down the line, we should make sure
that we are still abstracting away the Model from the end user. The viewer function
should probably be done by adding a new iterator on the base data manager that will
talk to the actual viewer. This is probably ideal to have the viewer be
primarily located in the data manager (first because it makes sense as the
viewers main job in this context is to feed in data for the model to load)
so that we can have an easier time ensuring that the viewer is always
returning the same formatted data as for in train / eval. All this is pending changes to
be done in the future... but just bear in mind that if learned parameters are in the data manager,
the viewer may have to use those parameters as well.
Args:
config: configuration to instantiate pipeline
device: location to place model and data
@ -94,11 +81,115 @@ class Pipeline(nn.Module):
self.model: The model that will be used
"""
datamanager: DataManager
model: Model
@property
def device(self):
"""Returns the device that the model is on."""
return module_wrapper(self.model).device
@profiler.time_function
def get_train_loss_dict(self, step: int):
"""This function gets your training loss dict. This will be responsible for
getting the next batch of data from the DataManager and interfacing with the
Model class, feeding the data to the model's forward function.
Args:
step: current iteration step to update sampler if using DDP (distributed)
"""
if self.world_size > 1 and step:
assert self.datamanager.train_sampler is not None
self.datamanager.train_sampler.set_epoch(step)
ray_bundle, batch = self.datamanager.next_train(step)
model_outputs = self.model(ray_bundle, batch)
metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)
return model_outputs, loss_dict, metrics_dict
@profiler.time_function
def get_eval_loss_dict(self, step: int):
"""This function gets your evaluation loss dict. It needs to get the data
from the DataManager and feed it to the model's forward function
Args:
step: current iteration step
"""
self.eval()
if self.world_size > 1:
assert self.datamanager.eval_sampler is not None
self.datamanager.eval_sampler.set_epoch(step)
ray_bundle, batch = self.datamanager.next_eval(step)
model_outputs = self.model(ray_bundle, batch)
loss_dict = self.model.get_loss_dict(model_outputs, batch)
metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
self.train()
return model_outputs, loss_dict, metrics_dict
@abstractmethod
@profiler.time_function
def get_eval_image_metrics_and_images(self, step: int):
"""This function gets your evaluation loss dict. It needs to get the data
from the DataManager and feed it to the model's forward function
Args:
step: current iteration step
"""
@abstractmethod
@profiler.time_function
def get_average_eval_image_metrics(self, step: Optional[int] = None):
"""Iterate over all the images in the eval dataset and get the average."""
@abstractmethod
def log_test_image_outputs(self) -> None:
"""Log the test image outputs"""
def load_pipeline(self, loaded_state: Dict[str, Any]) -> None:
"""Load the checkpoint from the given path
Args:
loaded_state: pre-trained model state dict
"""
def get_training_callbacks(
self, training_callback_attributes: TrainingCallbackAttributes
) -> List[TrainingCallback]:
"""Returns the training callbacks from both the Dataloader and the Model."""
def get_param_groups(self) -> Dict[str, List[Parameter]]:
"""Get the param groups for the pipeline.
Returns:
A list of dictionaries containing the pipeline's param groups.
"""
class VanillaPipeline(Pipeline):
"""The pipeline class for the vanilla nerf setup of multiple cameras for one or a few scenes.
config: configuration to instantiate pipeline
device: location to place model and data
test_mode: if True, loads test datset. if False, loads train/eval datasets
world_size: total number of machines available
local_rank: rank of current machine
Attributes:
self.datamanager: The data manager that will be used
self.model: The model that will be used
"""
def __init__(
self, config: cfg.PipelineConfig, device: str, test_mode: bool = False, world_size: int = 1, local_rank: int = 0
self,
config: cfg.VanillaPipelineConfig,
device: str,
test_mode: bool = False,
world_size: int = 1,
local_rank: int = 0,
):
super().__init__()
self.datamanager: DataManager = config.datamanager.setup(
self.datamanager: VanillaDataManager = config.datamanager.setup(
device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank
)
self.datamanager.to(device)