added caching for dataloader

mattport/nerf/dataset/collate.py

@@ -0,0 +1,62 @@
+"""
+Classes to help with caching data while training.
+"""
+
+from typing import Callable
+import torch
+import random
+
+from torch.utils.data import default_collate
+
+
+def collate_batch_size_one(batch_list):
+    assert len(batch_list) == 1
+    collated_batch = default_collate(batch_list)
+    for key in collated_batch:
+        assert collated_batch[key].shape[0] == 1
+        collated_batch[key] = collated_batch[key].squeeze(0)
+    return collated_batch
+
+
+class CollateIterDataset(torch.utils.data.IterableDataset):
+    """Collated image dataset that implements caching of images.
+
+    Args:
+        torch (_type_): _description_
+    """
+
+    def __init__(
+        self,
+        dataset: torch.utils.data.Dataset,
+        collate_fn: Callable,
+        num_samples_to_collate: int,
+        num_times_to_repeat: int = 0,
+    ):
+        super().__init__()
+        self.dataset = dataset
+        self.collate_fn = collate_fn
+        self.num_samples_to_collate = num_samples_to_collate
+        self.num_times_to_repeat = num_times_to_repeat
+
+        self.num_repeated = self.num_times_to_repeat  # starting value
+        self.cached_batch_list = None
+
+    def get_batch_list(self):
+        # sampling without replacement. TODO(ethan): don't use 'random'
+        indices = random.sample(range(len(self.dataset)), k=self.num_samples_to_collate)
+        batch_list = [self.dataset.__getitem__(idx) for idx in indices]
+        return batch_list
+
+    def __iter__(self):
+        while True:
+            if self.num_repeated >= self.num_times_to_repeat:
+                # trigger a reset
+                self.num_repeated = 0
+                batch_list = self.get_batch_list()
+                # possibly save a cached item
+                self.cached_batch_list = batch_list if self.num_times_to_repeat != 0 else None
+            else:
+                batch_list = self.cached_batch_list
+                self.num_repeated += 1
+            collated_batch = self.collate_fn(batch_list)
+            yield collated_batch

mattport/nerf/dataset/image_dataset.py

@@ -10,7 +10,38 @@ import PIL
 from PIL import Image
 from torch.utils.data import default_collate
 
-from mattport.utils.misc import DotDict
+
+def collate_batch(batch_list, num_rays_per_batch, keep_full_image: bool = False):
+    """_summary_
+
+    Args:
+     batch (_type_): _description_
+    """
+    # TODO(ethan): handle sampling even when in different processes
+    # we don't want the same samples for all spawned processed when
+    # using distributed training
+
+    batch = default_collate(batch_list)
+    num_images, image_height, image_width, _ = batch["image"].shape
+    indices = torch.floor(
+        torch.rand((num_rays_per_batch, 3)) * torch.tensor([num_images, image_height, image_width])
+    ).long()
+    c, y, x = [i.flatten() for i in torch.split(indices, 1, dim=-1)]
+    pixels = batch["image"][c, y, x]
+    mask = batch["mask"][c, y, x]
+    assert pixels.shape == (num_rays_per_batch, 3), pixels.shape
+
+    # Needed to correct the random indices to their actual camera idx locations.
+    indices[:, 0] = batch["image_idx"][c]
+    collated_batch = {
+        "indices": indices,
+        "pixels": pixels,
+        "mask": mask,
+    }
+    if keep_full_image:
+        collated_batch["image"] = batch["image"]
+
+    return collated_batch
 
 
 class ImageDataset(torch.utils.data.Dataset):
@@ -55,9 +86,6 @@ class ImageDataset(torch.utils.data.Dataset):
         return image
 
     def __getitem__(self, image_idx):
-        # TODO(ethan): we should add code to deal with varying resolution and/or having a maximum pixel resolution
-        # 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 = torch.from_numpy(self.get_image(image_idx).astype("float32") / 255.0)
         image = original_image[:, :, :3]
@@ -69,41 +97,9 @@ class ImageDataset(torch.utils.data.Dataset):
         else:
             raise ValueError(f"Image shape of {image.shape} is in correct.")
 
-        return {
+        data = {
             "image_idx": image_idx,
             "image": image,  # the pixels
             "mask": mask,
         }
-
-
-def collate_batch(batch, num_rays_per_batch, keep_full_image: bool = False):
-    """_summary_
-
-    Args:
-     batch (_type_): _description_
-    """
-    # TODO(ethan): handle sampling even when in different processes
-    # we don't want the same samples for all spawned processed when
-    # using distributed training
-
-    batch = default_collate(batch)
-    num_images, image_height, image_width, _ = batch["image"].shape
-    indices = torch.floor(
-        torch.rand((num_rays_per_batch, 3)) * torch.tensor([num_images, image_height, image_width])
-    ).long()
-    c, y, x = [i.flatten() for i in torch.split(indices, 1, dim=-1)]
-    pixels = batch["image"][c, y, x]
-    mask = batch["mask"][c, y, x]
-    assert pixels.shape == (num_rays_per_batch, 3), pixels.shape
-
-    indices[:, 0] = batch["image_idx"][c]
-
-    collated_batch = {
-        "indices": indices,
-        "pixels": pixels,
-        "mask": mask,
-    }
-    if keep_full_image:
-        collated_batch["image"] = batch["image"]
-
-    return DotDict(collated_batch)
+        return data

mattport/nerf/graph/vanilla_nerf.py

@@ -99,7 +99,10 @@ class NeRFGraph(Graph):
         self.sampler_uniform = UniformSampler(
             near_plane=self.near_plane, far_plane=self.far_plane, num_samples=self.num_coarse_samples
         )
-        self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples)
+        # self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples)
+        self.sampler_pdf = UniformSampler(
+            near_plane=self.near_plane, far_plane=self.far_plane, num_samples=self.num_importance_samples
+        )
 
         # field
         self.field_coarse = NeRFField()
@@ -128,7 +131,10 @@ class NeRFGraph(Graph):
         ray_bundle = self.ray_generator.forward(ray_indices)  # RayBundle
         # coarse network:
         uniform_ray_samples = self.sampler_uniform(ray_bundle)  # RaySamples
+
+        # time coarse network # xyz -> asd
         coarse_field_outputs = self.field_coarse(uniform_ray_samples)  # FieldOutputs
+        # time end coarse network
 
         coarse_renderer_outputs = self.renderer_rgb(
             rgb=coarse_field_outputs[FieldHeadNames.RGB],
@@ -136,7 +142,8 @@ class NeRFGraph(Graph):
             deltas=uniform_ray_samples.deltas,
         )  # RendererOutputs
         # fine network:
-        pdf_ray_samples = self.sampler_pdf(ray_bundle, uniform_ray_samples, coarse_field_outputs)  # RaySamples
+        # pdf_ray_samples = self.sampler_pdf(ray_bundle, uniform_ray_samples, coarse_field_outputs)  # RaySamples
+        pdf_ray_samples = self.sampler_pdf(ray_bundle)  # RaySamples
         fine_field_outputs = self.field_fine(pdf_ray_samples)  # FieldOutputs
 
         fine_renderer_outputs = self.renderer_rgb(
@@ -151,7 +158,7 @@ class NeRFGraph(Graph):
     def get_losses(self, batch, graph_outputs):
         # batch.pixels # (num_rays, 3)
         losses = {}
-        rgb_loss_coarse = self.rgb_loss(batch.pixels, graph_outputs["rgb_coarse"])
-        rgb_loss_fine = self.rgb_loss(batch.pixels, graph_outputs["rgb_fine"])
+        rgb_loss_coarse = self.rgb_loss(batch["pixels"], graph_outputs["rgb_coarse"])
+        rgb_loss_fine = self.rgb_loss(batch["pixels"], graph_outputs["rgb_fine"])
         losses = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine}
         return losses

mattport/nerf/trainer.py

@@ -14,6 +14,7 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader
 
 from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch
+from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one
 from mattport.nerf.dataset.utils import get_dataset_inputs
 from mattport.nerf.metrics import get_psnr
 from mattport.nerf.optimizers import Optimizers
@@ -50,20 +51,25 @@ class Trainer:
     def setup_dataset(self):
         """_summary_"""
         dataset_inputs = get_dataset_inputs(**self.config.dataset)
-        self.train_dataset = ImageDataset(
+        self.train_image_dataset = ImageDataset(
             image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor
         )
+        self.train_dataset = CollateIterDataset(
+            self.train_image_dataset,
+            collate_fn=lambda batch_list: collate_batch(
+                batch_list, self.config.dataloader.num_rays_per_batch, keep_full_image=False
+            ),
+            num_samples_to_collate=self.config.dataloader.num_images_to_sample_from,
+            num_times_to_repeat=self.config.dataloader.num_times_to_repeat_images,
+        )
         self.train_dataloader = DataLoader(
             self.train_dataset,
-            batch_size=self.config.dataloader.num_images_to_sample_from,
-            collate_fn=lambda batch: collate_batch(
-                batch, self.config.dataloader.num_rays_per_batch, keep_full_image=True
-            ),
+            batch_size=1,
             num_workers=self.config.dataloader.num_workers,
-            shuffle=True,
+            collate_fn=lambda batch_list: collate_batch_size_one(batch_list),
             pin_memory=True,
         )
-        # TODO(ethan): implement the test data
+        # TODO(ethan): implement the test dataset
 
     def setup_graph(self):
         """_summary_"""
@@ -175,13 +181,16 @@ class Trainer:
         """_summary_"""
         train_start = time()
         num_iterations = self.config.max_num_iterations
+        iter_dataset = iter(self.train_dataloader)
         for step in range(self.start_step, self.start_step + num_iterations):
             data_start = time()
-            batch = next(iter(self.train_dataloader))
+            batch = next(iter_dataset)
             self.update_time("load data (time)", data_start, time(), num_iter=step)
             iter_start = time()
             loss_dict = self.train_iteration(batch, step)
-            self.update_time("avg. rays/s (1/s)", iter_start, time(), num_iter=step, batch_size=batch.indices.shape[0])
+            self.update_time(
+                "avg. rays/s (1/s)", iter_start, time(), num_iter=step, batch_size=batch["indices"].shape[0]
+            )
             self.update_time("single iter (time)", iter_start, time(), num_iter=step)
             if step != 0 and step % self.config.steps_per_log == 0:
                 self.tensorboard_writer.write_scalar_dict(loss_dict, step, group="Loss", prefix="train-")
@@ -192,6 +201,8 @@ class Trainer:
                 logging.info("Saved ckpt at {%d}", step)
             if step != 0 and step % self.config.steps_per_test == 0:
                 self.test_image(image_idx=0, step=step)
+                self.test_image(image_idx=10, step=step)
+                self.test_image(image_idx=20, step=step)
             self.print_times()
         self.update_time("total train", train_start, time())
 
@@ -212,7 +223,7 @@ class Trainer:
     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 = self.train_image_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()
@@ -233,7 +244,9 @@ class Trainer:
             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)
+        self.tensorboard_writer.write_image(
+            f"image_idx_{image_idx}-rgb_coarse_fine", combined_image, step, group="test"
+        )
 
         fine_psnr = get_psnr(image, rgb_fine)
-        self.tensorboard_writer.write_scalar("fine_psnr", fine_psnr, step, group="test")
+        self.tensorboard_writer.write_scalar(f"image_idx_{image_idx}-fine_psnr", fine_psnr, step, group="test")

mattport/structures/cameras.py

@@ -181,10 +181,10 @@ class PinholeCamera(Camera):
         fy = intrinsics[:, cls.fy_index() : cls.fy_index() + 1]
         y = coords[:, 0:1]  # (num_rays, 1)
         x = coords[:, 1:2]  # (num_rays, 1)
-        directions = torch.cat([(x - cx) / fx, -(y - cy) / fy, -torch.ones_like(x)], -1)  # (num_rays, 3)
+        original_directions = torch.cat([(x - cx) / fx, -(y - cy) / fy, -torch.ones_like(x)], -1)  # (num_rays, 3)
         rotation = camera_to_world[:, :3, :3]  # (num_rays, 3, 3)
         directions = torch.sum(
-            directions[:, None, :] * rotation, dim=-1
+            original_directions[:, None, :] * rotation, dim=-1
         )  # (num_rays, 1, 3) * (num_rays, 3, 3) -> (num_rays, 3)
         directions = normalize(directions, dim=-1)
         origins = camera_to_world[:, :3, 3]  # (num_rays, 3)