nerfstudio/notebooks/creating_pipelines.ipynb

Creating pipelines

Here we explain how to create custom pipelines in nerfactory. Pipelines are composed of two components, namely a Dataloader and a Model.

In [5]:

# HIDDEN
%load_ext autoreload
%autoreload 2

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

In [6]:

from typing import Dict
import torch

from nerfactory.dataloaders.base import Dataloader
from nerfactory.models.base import Model
from nerfactory.pipelines.base import Pipeline

from nerfactory.models.instant_ngp import NGPModel
from nerfactory.cameras.rays import RayBundle


class SemanticNGPModel(Model):
    """An instant ngp model modified slightly to output semantics."""

    def __init__(self):
        self.ngp_model = NGPModel()

    def get_outputs(self, ray_bundle: RayBundle) -> Dict[str, torch.Tensor]:
        outputs = self.ngp_model.forward(ray_bundle)
        outputs["semantics"] = torch.rand_like(outputs["rgb"])
        return outputs


class CustomPipeline(Pipeline):
    """The Instant NGP pipeline."""

    dataloader: Dataloader
    model: Model

    def get_train_loss_dict(self):
        rays, batch = self.dataloader_train_iter.next()
        # TODO: do something here
        accumulated_color, _, _, mask = self.model(rays, batch)
        masked_batch = get_masked_dict(batch, mask)
        loss_dict = self.model.get_loss_dict(accumulated_color, masked_batch, mask)
        return loss_dict

    def get_eval_loss_dict(self):
        rays, batch = self.dataloader_eval_iter.next()
        # TODO: do something here
        accumulated_color, _, _, mask = self.model(rays, batch)
        masked_batch = get_masked_dict(batch, mask)
        loss_dict = self.model.get_loss_dict(accumulated_color, masked_batch, mask)
        return loss_dict

    def test_image_outputs(self) -> None:
        """Log the test image outputs"""
        # camera = Camera(self.dataloader.eval_datasetinputs.camera_to_world[0])
        pass

In [9]:

dataloader = Dataloader(use_train=True, use_eval=True)
model = SemanticNGPModel()
pipeline = CustomPipeline(dataloader=dataloader, model=model)

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
/shared/ethanweber/nerfactory/notebooks/creating_pipelines.ipynb Cell 4 in <cell line: 1>()
----> <a href='vscode-notebook-cell://ssh-remote%2Bbirman/shared/ethanweber/nerfactory/notebooks/creating_pipelines.ipynb#W3sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0'>1</a> dataloader = Dataloader(use_train=True, use_eval=True)
      <a href='vscode-notebook-cell://ssh-remote%2Bbirman/shared/ethanweber/nerfactory/notebooks/creating_pipelines.ipynb#W3sdnNjb2RlLXJlbW90ZQ%3D%3D?line=1'>2</a> model = SemanticNGPModel()
      <a href='vscode-notebook-cell://ssh-remote%2Bbirman/shared/ethanweber/nerfactory/notebooks/creating_pipelines.ipynb#W3sdnNjb2RlLXJlbW90ZQ%3D%3D?line=2'>3</a> pipeline = CustomPipeline(dataloader=dataloader, model=model)

File /shared/ethanweber/nerfactory/nerfactory/dataloaders/base.py:95, in Dataloader.__init__(self, use_train, use_eval)
     93 assert use_train or use_eval
     94 if use_train:
---> 95     self.setup_train()
     96 if use_eval:
     97     self.setup_eval()

File /shared/ethanweber/nerfactory/nerfactory/dataloaders/base.py:116, in Dataloader.setup_train(self)
    113 @abstractmethod
    114 def setup_train(self):
    115     """Sets up the dataloader for training"""
--> 116     raise NotImplementedError

NotImplementedError: 

Creating pipelines from a config

Now we show how to create a pipeline from a config, which has the following form:

@dataclass
class PipelineConfig:
    """Configuration for pipeline instantiation"""

    _target_: str = MISSING
    dataloader_config: DataloaderConfig = MISSING
    graph_config: GraphConfig = MISSING

See nerfactory/utils/config.py for more details. In this example, we will simply load from an existing configuration from configs/graph_instant_ngp.yaml.

In [10]:

import pprint
import hydra

hydra.core.global_hydra.GlobalHydra.instance().clear()
from hydra import compose, initialize

initialize(version_base="1.2", config_path="../configs/")
config_name = "graph_instant_ngp.yaml"
config = compose(config_name)
pipeline_config = config.pipeline
pprint.pprint(pipeline_config)
print("----------------------------------------------------")

from nerfactory.pipelines.base import setup_pipeline

pipeline = setup_pipeline(pipeline_config, device="cuda")

{'_target_': 'nerfactory.pipelines.base.Pipeline', 'dataloader': {'_target_': 'nerfactory.dataloaders.base.VanillaDataloader', 'train_dataset': {'_target_': 'nerfactory.dataloaders.datasets.Blender', 'data_directory': 'data/blender/lego', 'alpha_color': 'white', 'downscale_factor': 1}, 'train_num_rays_per_batch': 8192, 'eval_dataset': {'_target_': 'nerfactory.dataloaders.datasets.Blender', 'data_directory': 'data/blender/lego', 'alpha_color': 'white', 'downscale_factor': 1}, 'eval_num_rays_per_chunk': 4096, 'eval_num_rays_per_chunk': 8192}, 'model': {'_target_': 'nerfactory.models.instant_ngp.NGPModel', 'enable_collider': False, 'collider_config': {'_target_': 'nerfactory.models.modules.scene_colliders.NearFarCollider', 'near_plane': 2.0, 'far_plane': 6.0}, 'num_coarse_samples': 64, 'num_importance_samples': 128, 'loss_coefficients': {'rgb_loss_coarse': 1.0, 'rgb_loss_fine': 1.0, 'rgb_loss': 1.0}, 'enable_density_field': True, 'density_field_config': {'_target_': 'nerfactory.fields.density_fields.density_grid.DensityGrid', 'center': 0.0, 'base_scale': 3.0, 'num_cascades': 1, 'resolution': 128, 'update_every_num_iters': 16}, 'field_implementation': 'tcnn'}}
----------------------------------------------------

/shared/ethanweber/miniconda3/envs/nerfactory/lib/python3.8/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and will be removed in 0.15, please use 'weights' instead.
  warnings.warn(
/shared/ethanweber/miniconda3/envs/nerfactory/lib/python3.8/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and will be removed in 0.15. The current behavior is equivalent to passing `weights=AlexNet_Weights.IMAGENET1K_V1`. You can also use `weights=AlexNet_Weights.DEFAULT` to get the most up-to-date weights.
  warnings.warn(msg)

In [11]:

pipeline.get_train_loss_dict()

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/shared/ethanweber/nerfactory/notebooks/creating_pipelines.ipynb Cell 8 in <cell line: 1>()
----> <a href='vscode-notebook-cell://ssh-remote%2Bbirman/shared/ethanweber/nerfactory/notebooks/creating_pipelines.ipynb#X13sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0'>1</a> pipeline.get_train_loss_dict()

File /shared/ethanweber/nerfactory/nerfactory/pipelines/base.py:81, in Pipeline.get_train_loss_dict(self)
     76 @abstractmethod
     77 def get_train_loss_dict(self):
     78     """This function gets your training loss dict. This will be responsible for
     79     getting the next batch of data from the dataloader and interfacing with the
     80     Model class, feeding the data to the model's forward function."""
---> 81     rays, batch = self.dataloader_train_iter.next()
     82     accumulated_color, _, _, mask = self.model(rays, batch)
     83     masked_batch = get_masked_dict(batch, mask)

AttributeError: 'IterableWrapper' object has no attribute 'next'

In [ ]: