diff --git a/mattport/nerf/dataset/collate.py b/mattport/nerf/dataset/collate.py new file mode 100644 index 000000000..54adb164c --- /dev/null +++ b/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 diff --git a/mattport/nerf/dataset/image_dataset.py b/mattport/nerf/dataset/image_dataset.py index 0acdc2405..dd9eee767 100644 --- a/mattport/nerf/dataset/image_dataset.py +++ b/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 diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 7696eb237..8e3320e16 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/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 diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 3585b02f9..51349697a 100644 --- a/mattport/nerf/trainer.py +++ b/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") diff --git a/mattport/structures/cameras.py b/mattport/structures/cameras.py index 0eca34707..6e8f6e5e1 100644 --- a/mattport/structures/cameras.py +++ b/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)