added test set (using training data for now)

.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,12 +1,21 @@
+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
 
 max_num_iterations: 1000000
-steps_per_save: 1000
-steps_per_log: 100
-steps_per_test: 5000
+steps_per_save: null
+steps_per_log: 10
+steps_per_test: 100
 
 model_dir: /tmp/mattport_models/
 
@@ -25,7 +34,7 @@ dataloader:
 
 # optimizer options
 param_groups:
-  graph:
+  fields:
     optimizer:
       type: Adam
       lr: 1e-4
@@ -33,7 +42,7 @@ param_groups:
       type: StepLR
       step_size: 30
       gamma: 0.1
-  camera:
+  cameras:
     optimizer:
       type: SGD
       lr: 1e-4
@@ -43,7 +52,7 @@ param_groups:
 
 # machine options
 machine_config:
-  num_gpus: 4
+  num_gpus: 1
   num_machines: 1
   machine_rank: 0
   dist_url: auto

mattport/nerf/dataset/image_dataset.py

@@ -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/graph/nerf.py

@@ -51,8 +51,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]):

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

@@ -2,8 +2,6 @@
 Optimizers class.
 """
 
-import logging
-
 from typing import Dict, List
 from omegaconf import DictConfig
 import torch
@@ -53,16 +51,16 @@ class Optimizers:
 
     def zero_grad_all(self):
         """zero the gradients for all optimizer parameters"""
-        for param_name, optimizer in self.optimizers.items():
-            # print(param_name)
-            pass
+        for _, optimizer in self.optimizers.items():
+            optimizer.zero_grad()
 
     def optimizer_step_all(self):
         """Run step for all optimizers."""
         # logging.info("optimizer_step_all")
-        pass
+        for param_name, optimizer in self.optimizers.items():
+            optimizer.step()
 
-    def scheduler_step_all(self):
+    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_"""
@@ -118,40 +122,69 @@ class Trainer:
         )
 
     def get_aggregated_loss(self, losses):
-        aggregated_loss = 0.0
+        loss_sum = 0.0
+        loss_dict = {}
         for loss_name in losses.keys():
             # TODO(ethan): add loss weightings here from a config
-            aggregated_loss += losses[loss_name]
-        return aggregated_loss
+            loss_sum += losses[loss_name]
+            loss_dict[loss_name] = float(losses[loss_name])
+        return loss_sum, loss_dict
 
     def train(self) -> None:
         """_summary_"""
-        num_iterations = 1000
+        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))
-            losses = self.train_iteration(step, batch)
-            self.local_writer.write_text(f"{step}/{num_iterations}: losses", losses)
-            if step % self.config.steps_per_save == 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, step: int, batch: dict):
+    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}")
-        # print(batch)
         losses = self.graph.get_losses(batch, graph_outputs)  # or self.graph.module
-        loss = self.get_aggregated_loss(losses)
+        loss_sum, loss_dict = self.get_aggregated_loss(losses)
         self.optimizers.zero_grad_all()
-        loss.backward()
-        self.optimizers.scheduler_step_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 losses
+        return loss_dict
 
-    def test(self, image_idx):
+    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():
-            ray_indices = None
-            pass
+            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/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

@@ -179,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(