Merge pull request #35 from ethanweber/ethan/test_set_integration

Fake version of test set with tensorboard.

.vscode/launch.json

@@ -8,9 +8,11 @@
             "name": "Python: run_train_nerf",
             "type": "python",
             "request": "launch",
-            "program": "scripts/run_train_nerf.py",
+            "program": "scripts/run_train.py",
             "console": "integratedTerminal",
-            "args": []
+            "args": [
+                "machine_config.num_gpus=1"
+            ]
         }
     ]
 }

README.md

@@ -37,7 +37,10 @@ Download the original NeRF dataset and put it in the following format.
 # Training a model
 
 ```
-python scripts/run_train_nerf.py
+# Run with the default config
+python scripts/run_train.py
+# Run with config changes
+python scripts/run_train.py machine_config.num_gpus=1
 ```
 
 # Setting up Jupyter

configs/default.yaml

@@ -1,9 +1,22 @@
+hydra: 
+  run:
+    # Configure output dir of each experiment programmatically from the arguments
+    # Example "outputs/mnist/classifier/baseline/2021-03-10-141516"
+    dir: outputs/${experiment_name}/${method_name}/${now:%Y-%m-%d_%H%M%S}
+
+experiment_name: blender_lego
+method_name: vanilla_nerf
+
 # pretrained loading options
 resume_train:
   load_dir: # pretrained model directory
   load_step: # pretrained model's step number
 
-save_step: 10
+max_num_iterations: 1000000
+steps_per_save: null
+steps_per_log: 10
+steps_per_test: 100
+
 model_dir: /tmp/mattport_models/
 
 # dataset options
@@ -21,7 +34,7 @@ dataloader:
 
 # optimizer options
 param_groups:
-  graph:
+  fields:
     optimizer:
       type: Adam
       lr: 1e-4
@@ -29,18 +42,17 @@ param_groups:
       type: StepLR
       step_size: 30
       gamma: 0.1
-  camera:
+  cameras:
     optimizer:
       type: SGD
       lr: 1e-4
     scheduler:
-      type: ExponentialLR 
+      type: ExponentialLR
       gamma: 0.1
 
-
 # machine options
 machine_config:
-  num_gpus: 4
+  num_gpus: 1
   num_machines: 1
   machine_rank: 0
   dist_url: auto
@@ -48,4 +60,4 @@ machine_config:
 # graph options
 graph:
   # NOTE(ethan): currently setting to nothing b/c we're still designing the config format...
-  _target_: mattport.nerf.graph.nerf.NeRFGraph
+  _target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph

mattport/nerf/dataset/image_dataset.py

@@ -9,7 +9,7 @@ import torch
 from PIL import Image
 from torch.utils.data import default_collate
 
-from mattport.nerf.dataset.utils import DotDict
+from mattport.utils.misc import DotDict
 
 
 class ImageDataset(torch.utils.data.Dataset):
@@ -43,7 +43,7 @@ class ImageDataset(torch.utils.data.Dataset):
         # e.g., we should be able to downscale the images or set a max height, width, etc.
 
         # the image might be RGB or RGBA, so we separate it
-        original_image = self.get_image(image_idx).astype("float32") / 255.0
+        original_image = torch.from_numpy(self.get_image(image_idx).astype("float32") / 255.0)
         image = original_image[:, :, :3]
         num_channels = original_image.shape[2]
         if num_channels == 4:

mattport/nerf/dataset/utils.py

@@ -13,18 +13,6 @@ import torch
 from mattport.utils.io import get_absolute_path, load_from_json
 
 
-class DotDict(dict):
-    """
-    dot.notation access to dictionary attributes
-    """
-
-    def __getattr__(self, attr):
-        return self[attr]
-
-    __setattr__ = dict.__setitem__
-    __delattr__ = dict.__delitem__
-
-
 @dataclass
 class DatasetInputs:
     """Dataset inputs for the image dataset and the ray generator."""

mattport/nerf/graph/vanilla_nerf.py

@@ -11,7 +11,7 @@ from torch.nn import Parameter
 from torchtyping import TensorType
 
 from mattport.nerf.field_modules.encoding import NeRFEncoding
-from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead
+from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHeadNames, RGBFieldHead
 from mattport.nerf.field_modules.mlp import MLP
 from mattport.nerf.field_modules.ray_generator import RayGenerator
 from mattport.nerf.graph.base import Graph
@@ -58,7 +58,9 @@ class NeRFField(nn.Module):
         field_rgb_output = self.field_output_rgb(rgb_mlp_out)
         field_density_out = self.field_output_density(base_mlp_out)
 
-        field_outputs = field_rgb_output | field_density_out
+        field_outputs = {}
+        field_outputs.update(field_rgb_output)
+        field_outputs.update(field_density_out)
         return field_outputs
 
 
@@ -73,7 +75,7 @@ class NeRFGraph(Graph):
         self.ray_generator = RayGenerator(self.intrinsics, self.camera_to_world)
 
         # samplers
-        self.sampler_uniform = UniformSampler(near_plane=0.1, far_plane=4.0, num_samples=64)
+        self.sampler_uniform = UniformSampler(near_plane=0.1, far_plane=6.0, num_samples=64)
         self.sampler_pdf = PDFSampler(num_samples=64)
 
         # field
@@ -93,8 +95,8 @@ class NeRFGraph(Graph):
             Dict[str, List[Parameter]]: Mapping of different parameter groups
         """
         param_groups = {}
-        param_groups["camera"] = [dict(params=self.ray_generator.parameters())]
-        param_groups["graph"] = [dict(params=self.field_coarse.parameters()), dict(params=self.field_fine.parameters())]
+        param_groups["cameras"] = list(self.ray_generator.parameters())
+        param_groups["fields"] = list(self.field_coarse.parameters()) + list(self.field_fine.parameters())
         return param_groups
 
     def forward(self, ray_indices: TensorType["num_rays", 3]):
@@ -104,14 +106,20 @@ class NeRFGraph(Graph):
         # coarse network:
         uniform_ray_samples = self.sampler_uniform(ray_bundle)  # RaySamples
         coarse_field_outputs = self.field_coarse(uniform_ray_samples)  # FieldOutputs
+
         coarse_renderer_outputs = self.renderer_rgb(
-            field_outputs=coarse_field_outputs.rgb, deltas=uniform_ray_samples.deltas
+            rgb=coarse_field_outputs[FieldHeadNames.RGB],
+            density=coarse_field_outputs[FieldHeadNames.DENSITY],
+            deltas=uniform_ray_samples.deltas,
         )  # RendererOutputs
         # fine network:
         pdf_ray_samples = self.sampler_pdf(ray_bundle, uniform_ray_samples, coarse_field_outputs)  # RaySamples
         fine_field_outputs = self.field_fine(pdf_ray_samples)  # FieldOutputs
+
         fine_renderer_outputs = self.renderer_rgb(
-            field_outputs=fine_field_outputs, deltas=pdf_ray_samples.deltas
+            rgb=fine_field_outputs[FieldHeadNames.RGB],
+            density=fine_field_outputs[FieldHeadNames.DENSITY],
+            deltas=pdf_ray_samples.deltas,
         )  # RendererOutputs
         # outputs:
         outputs = {"rgb_coarse": coarse_renderer_outputs.rgb, "rgb_fine": fine_renderer_outputs.rgb}

mattport/nerf/metrics.py

@@ -0,0 +1,13 @@
+"""
+Collection of metrics.
+"""
+
+import torch
+
+
+def get_psnr(image_gt, image_pred):
+    """
+    Returns the psnr for an image. TODO: add masking
+    """
+    mean = torch.mean((image_gt - image_pred) ** 2)
+    return -10 * torch.log10(mean)

mattport/nerf/optimizers.py

@@ -48,3 +48,19 @@ class Optimizers:
         """
         if self.config.param_name.scheduler:
             self.schedulers[param_name].step()
+
+    def zero_grad_all(self):
+        """zero the gradients for all optimizer parameters"""
+        for _, optimizer in self.optimizers.items():
+            optimizer.zero_grad()
+
+    def optimizer_step_all(self):
+        """Run step for all optimizers."""
+        # logging.info("optimizer_step_all")
+        for param_name, optimizer in self.optimizers.items():
+            optimizer.step()
+
+    def scheduler_step_all(self, step):
+        """Run step for all schedulers."""
+        # logging.info("optimizer_step_all")
+        pass

mattport/nerf/trainer.py

@@ -15,7 +15,8 @@ from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch
 from mattport.nerf.dataset.utils import get_dataset_inputs
 from mattport.nerf.optimizers import Optimizers
 from mattport.utils.decorators import check_main_thread
-from mattport.utils.writer import LocalWriter
+from mattport.utils.writer import LocalWriter, TensorboardWriter
+from mattport.nerf.metrics import get_psnr
 
 logging.getLogger("PIL").setLevel(logging.WARNING)
 
@@ -38,7 +39,10 @@ class Trainer:
         self.start_step = 0
         # logging variables
         self.is_main_thread = local_rank % world_size == 0
-        self.local_writer = LocalWriter(local_rank, world_size, save_dir=None)
+        self.local_writer = LocalWriter(local_rank, world_size, save_dir=os.path.join(os.getcwd(), "writer"))
+        self.tensorboard_writer = TensorboardWriter(
+            local_rank, world_size, save_dir=os.path.join(os.getcwd(), "writer")
+        )
 
     def setup_dataset(self):
         """_summary_"""
@@ -74,6 +78,7 @@ class Trainer:
 
     def setup_optimizers(self):
         """_summary_"""
+        # TODO(ethan): handle different world sizes
         self.optimizers = Optimizers(self.config.param_groups, self.graph.get_param_groups())
 
     def load_checkpoint(self, load_config: DictConfig) -> int:
@@ -116,23 +121,70 @@ class Trainer:
             ckpt_path,
         )
 
+    def get_aggregated_loss(self, losses):
+        loss_sum = 0.0
+        loss_dict = {}
+        for loss_name in losses.keys():
+            # TODO(ethan): add loss weightings here from a config
+            loss_sum += losses[loss_name]
+            loss_dict[loss_name] = float(losses[loss_name])
+        return loss_sum, loss_dict
+
     def train(self) -> None:
         """_summary_"""
-        num_epochs = 10
-        for _ in range(num_epochs):
-            self.train_iteration()
-
-    def train_iteration(self):
-        """_summary_"""
-        num_iters = 100
-        for step in range(self.start_step, self.start_step + num_iters):
+        num_iterations = self.config.max_num_iterations
+        for step in range(self.start_step, self.start_step + num_iterations):
             batch = next(iter(self.train_dataloader))
-            # move batch to correct device
-            ray_indices = batch.indices.to(f"cuda:{self.local_rank}")
-            graph_outputs = self.graph(ray_indices)
-            batch.pixels = batch.pixels.to(f"cuda:{self.local_rank}")
-            losses = self.graph.module.get_losses(batch, graph_outputs)
-            self.local_writer.write_text(f"{step}/{num_iters}: losses", losses)
-            if step % self.config.save_step == 0:
+            loss_dict = self.train_iteration(batch, step)
+            if step != 0 and step % self.config.steps_per_log == 0:
+                self.tensorboard_writer.write_scalar_dict(loss_dict, step, group="train")
+                logging.info(f"{step}/{num_iterations} with losses {loss_dict}")
+                # TODO: add the learning rates to tensorboard/logging
+            if self.config.steps_per_save and step != 0 and step % self.config.steps_per_save == 0:
                 self.save_checkpoint(self.config.model_dir, step)
-                self.local_writer.write_text(f"ckpt at {step}", f"{losses}\n")
+                logging.info(f"Saved ckpt at {step}")
+            if step != 0 and step % self.config.steps_per_test == 0:
+                self.test_image(image_idx=0, step=step)
+
+    def train_iteration(self, batch: dict, step: int):
+        """Run one iteration with a batch of inputs."""
+        # move batch to correct device
+        ray_indices = batch.indices.to(f"cuda:{self.local_rank}")
+        graph_outputs = self.graph(ray_indices)
+        batch.pixels = batch.pixels.to(f"cuda:{self.local_rank}")
+        losses = self.graph.get_losses(batch, graph_outputs)  # or self.graph.module
+        loss_sum, loss_dict = self.get_aggregated_loss(losses)
+        self.optimizers.zero_grad_all()
+        loss_sum.backward()
+        self.optimizers.scheduler_step_all(step)  # NOTE(ethan): I think the scheduler needs to know what step we are on
+        self.optimizers.optimizer_step_all()
+        return loss_dict
+
+    def test_image(self, image_idx, step):
+        """Test a specific image."""
+        logging.info("Running test iteration.")
+        image = self.train_dataset[image_idx]["image"]  # ground truth
+        image_height, image_width, _ = image.shape
+        pixel_coords = torch.meshgrid(torch.arange(image_height), torch.arange(image_width), indexing="ij")
+        pixel_coords = torch.stack(pixel_coords, dim=-1).long()
+        all_ray_indices = torch.cat([torch.ones_like(pixel_coords[..., :1]) * image_idx, pixel_coords], dim=-1).view(
+            -1, 3
+        )
+        with torch.no_grad():
+            num_rays = all_ray_indices.shape[0]
+            chunk_size = 1024
+            rgb_coarse = []
+            rgb_fine = []
+            for i in range(0, num_rays, chunk_size):
+                ray_indices = all_ray_indices[i : i + chunk_size].to(f"cuda:{self.local_rank}")
+                graph_outputs = self.graph(ray_indices)
+                rgb_coarse.append(graph_outputs["rgb_coarse"])
+                rgb_fine.append(graph_outputs["rgb_fine"])
+            rgb_coarse = torch.cat(rgb_coarse).view(image_height, image_width, 3).detach().cpu()
+            rgb_fine = torch.cat(rgb_fine).view(image_height, image_width, 3).detach().cpu()
+
+        combined_image = torch.cat([image, rgb_coarse, rgb_fine], dim=1)
+        self.tensorboard_writer.write_image("test/rgb_coarse_fine", combined_image, step)
+
+        fine_psnr = get_psnr(image, rgb_fine)
+        self.tensorboard_writer.write_scalar("fine_psnr", fine_psnr, step, group="test")

mattport/utils/misc.py

@@ -0,0 +1,10 @@
+class DotDict(dict):
+    """
+    dot.notation access to dictionary attributes
+    """
+
+    def __getattr__(self, attr):
+        return self[attr]
+
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__

mattport/utils/writer.py

@@ -5,6 +5,7 @@ Generic Writer class
 
 import os
 from abc import abstractmethod
+from typing import Dict
 
 import imageio
 import numpy as np
@@ -14,7 +15,7 @@ from torchtyping import TensorType
 
 from mattport.utils.decorators import check_main_thread
 
-to8b = lambda x: (255 * torch.clip(x, 0, 1)).to(torch.uint8)
+to8b = lambda x: (255 * torch.clamp(x, min=0, max=1)).to(torch.uint8)
 
 
 class Writer:
@@ -22,11 +23,10 @@ class Writer:
 
     def __init__(self, local_rank: int, world_size: int, save_dir: str):
         self.save_dir = save_dir
-        self.writer = None
         self.is_main_thread = local_rank % world_size == 0
 
     @abstractmethod
-    def write_image(self, name: str, x: TensorType["H", "W", 3]) -> None:
+    def write_image(self, name: str, x: TensorType["H", "W", 3], step: int) -> None:
         """_summary_
 
         Args:
@@ -46,7 +46,7 @@ class Writer:
         raise NotImplementedError
 
     @abstractmethod
-    def write_scalar(self, name: str, x: float, y: float) -> None:
+    def write_scalar(self, name: str, scalar: float, step: float, group: str = None) -> None:
         """Required method to write a single scalar value to the logger
 
         Args:
@@ -56,15 +56,15 @@ class Writer:
         """
         raise NotImplementedError
 
-    def write_scalar_dict(self, scalar_dict: dict) -> None:
+    def write_scalar_dict(self, scalar_dict: Dict[str, float], step: int, group: str = None) -> None:
         """Function that writes out all scalars from a given dictionary to the logger
 
         Args:
             scalar_dict (dict): dictionary containing all scalar values
         """
         if self.is_main_thread:
-            for name, (x, y) in scalar_dict.items():
-                self.writer.write_scalar(name, x, y)
+            for name, scalar in scalar_dict.items():
+                self.write_scalar(name, scalar, step, group=group)
 
 
 class TensorboardWriter(Writer):
@@ -73,10 +73,10 @@ class TensorboardWriter(Writer):
     def __init__(self, local_rank: int, world_size: int, save_dir: str):
         super().__init__(local_rank, world_size, save_dir)
         if self.is_main_thread:
-            self.writer = SummaryWriter(log_dir=self.save_dir)
+            self.tb_writer = SummaryWriter(log_dir=self.save_dir)
 
     @check_main_thread
-    def write_image(self, name: str, x: TensorType["H", "W", 3]) -> None:
+    def write_image(self, name: str, x: TensorType["H", "W", 3], step: int) -> None:
         """_summary_
 
         Args:
@@ -84,7 +84,7 @@ class TensorboardWriter(Writer):
             x (TensorType["H", "W", 3]): rendered image to write
         """
         x = to8b(x)
-        self.writer.add_images(name, x)
+        self.tb_writer.add_image(name, x, step, dataformats="HWC")
 
     @check_main_thread
     def write_text(self, name: str, x: str) -> None:
@@ -94,28 +94,25 @@ class TensorboardWriter(Writer):
             name (str): data identifier
             x (str): string to write
         """
-        self.writer.add_text(name, x)
+        self.tb_writer.add_text(name, x)
 
     @check_main_thread
-    def write_scalar(self, name: str, x: float, y: float) -> None:
+    def write_scalar(self, name: str, scalar: float, step: float, group: str = None) -> None:
         """Tensorboard method to write a single scalar value to the logger
 
         Args:
             name (str): data identifier
             x (float): x value to write
             y (float): y value to write
+            group (str)): a prefix to group tensorboard scalars
         """
-        self.writer.add_scalar(name, x, y)
+        prefix = f"{group}/" if group else ""
+        self.tb_writer.add_scalar(f"{prefix}{name}", scalar, step)
 
 
 class LocalWriter(Writer):
     """Local Writer Class"""
 
-    def __init__(self, local_rank: int, world_size: int, save_dir: str):
-        super().__init__(local_rank, world_size, save_dir)
-        if self.is_main_thread:
-            self.writer = SummaryWriter(log_dir=self.save_dir)
-
     @check_main_thread
     def write_image(self, name: str, x: TensorType["H", "W", 3]) -> None:
         """Logs image to a jpg file in save_directory

scripts/run_train.py

@@ -21,6 +21,9 @@ logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=loggin
 
 DEFAULT_TIMEOUT = timedelta(minutes=30)
 
+# speedup for when input size to model doesn't change (much)
+torch.backends.cudnn.benchmark = True
+
 
 def _find_free_port() -> str:
     """Find free port on open socket"""
@@ -176,7 +179,7 @@ def launch(
         main_func(local_rank=0, world_size=1, config=config)
 
 
-@hydra.main(config_path="../configs", config_name="config.yaml")
+@hydra.main(config_path="../configs", config_name="default.yaml")
 def main(config: DictConfig):
     """Main function."""
     launch(