From 80ad17700c28aa70b54d88ccca931d9e7ea8e9a6 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Sat, 23 Apr 2022 19:45:53 +0000 Subject: [PATCH] pylint and test cases for dataloader --- mattport/nerf/dataset/collate.py | 9 +++- mattport/nerf/dataset/image_dataset.py | 4 +- mattport/nerf/field_modules/mlp.py | 5 +- mattport/nerf/graph/base.py | 2 +- mattport/nerf/graph/vanilla_nerf.py | 7 +-- mattport/nerf/optimizers.py | 7 ++- mattport/nerf/trainer.py | 16 +++--- mattport/utils/misc.py | 5 ++ notebooks/visualize_dataset.ipynb | 60 +++++++++++---------- tests/nerf/dataloader/test_ray_generator.py | 48 ++++++++++------- 10 files changed, 94 insertions(+), 69 deletions(-) diff --git a/mattport/nerf/dataset/collate.py b/mattport/nerf/dataset/collate.py index 54adb164c..8e241d1e4 100644 --- a/mattport/nerf/dataset/collate.py +++ b/mattport/nerf/dataset/collate.py @@ -2,14 +2,15 @@ Classes to help with caching data while training. """ -from typing import Callable -import torch import random +from typing import Callable +import torch from torch.utils.data import default_collate def collate_batch_size_one(batch_list): + """Use the default collate function but then squeze to avoid the batch dimension being added.""" assert len(batch_list) == 1 collated_batch = default_collate(batch_list) for key in collated_batch: @@ -41,7 +42,11 @@ class CollateIterDataset(torch.utils.data.IterableDataset): self.num_repeated = self.num_times_to_repeat # starting value self.cached_batch_list = None + def __getitem__(self, idx): + return self.dataset.__getitem__(idx) + def get_batch_list(self): + """Returns a list of batches from the dataset attribute.""" # 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] diff --git a/mattport/nerf/dataset/image_dataset.py b/mattport/nerf/dataset/image_dataset.py index dd9eee767..89558f5ab 100644 --- a/mattport/nerf/dataset/image_dataset.py +++ b/mattport/nerf/dataset/image_dataset.py @@ -32,9 +32,11 @@ def collate_batch(batch_list, num_rays_per_batch, keep_full_image: bool = False) assert pixels.shape == (num_rays_per_batch, 3), pixels.shape # Needed to correct the random indices to their actual camera idx locations. + local_indices = indices.clone() indices[:, 0] = batch["image_idx"][c] collated_batch = { - "indices": indices, + "local_indices": local_indices, # local to the batch returned + "indices": indices, # with the abs camera indices "pixels": pixels, "mask": mask, } diff --git a/mattport/nerf/field_modules/mlp.py b/mattport/nerf/field_modules/mlp.py index 7fcce4c7f..d081f117d 100644 --- a/mattport/nerf/field_modules/mlp.py +++ b/mattport/nerf/field_modules/mlp.py @@ -1,8 +1,7 @@ """ Multi Layer Perceptron """ -from tkinter import X -from typing import List, Optional +from typing import Tuple, Optional import torch from torch import nn @@ -20,7 +19,7 @@ class MLP(FieldModule): out_dim: int, num_layers: int, layer_width: int = None, - skip_connections: List[int] = [], + skip_connections: Tuple[int] = (), activation: Optional[nn.Module] = nn.ReLU(), out_activation: Optional[nn.Module] = None, ) -> None: diff --git a/mattport/nerf/graph/base.py b/mattport/nerf/graph/base.py index 041b3e8b1..2bce70e44 100644 --- a/mattport/nerf/graph/base.py +++ b/mattport/nerf/graph/base.py @@ -34,7 +34,7 @@ class Node: class Graph(nn.Module): """_summary_""" - def __init__(self, intrinsics=None, camera_to_world=None, **kwargs) -> None: + def __init__(self, intrinsics=None, camera_to_world=None) -> None: super().__init__() self.intrinsics = intrinsics self.camera_to_world = camera_to_world diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 8e3320e16..61cee7b18 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -3,7 +3,7 @@ Implementation of vanilla nerf. """ -from typing import Dict, List +from typing import Dict, List, Tuple import torch from torch import nn @@ -17,14 +17,14 @@ from mattport.nerf.field_modules.ray_generator import RayGenerator from mattport.nerf.graph.base import Graph from mattport.nerf.loss import MSELoss from mattport.nerf.renderers import RGBRenderer -from mattport.nerf.sampler import PDFSampler, UniformSampler +from mattport.nerf.sampler import PDFSampler, UniformSampler # pylint: disable=unused-import from mattport.structures.rays import RaySamples class NeRFField(nn.Module): """NeRF module""" - def __init__(self, num_layers=8, layer_width=256, skip_connections=[4]) -> None: + def __init__(self, num_layers=8, layer_width=256, skip_connections: Tuple = (4)) -> None: super().__init__() self.num_layers = num_layers self.layer_width = layer_width @@ -99,6 +99,7 @@ class NeRFGraph(Graph): self.sampler_uniform = UniformSampler( near_plane=self.near_plane, far_plane=self.far_plane, num_samples=self.num_coarse_samples ) + # TODO: this needs to be corrected # 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 diff --git a/mattport/nerf/optimizers.py b/mattport/nerf/optimizers.py index 4431e9f51..3acb05a91 100644 --- a/mattport/nerf/optimizers.py +++ b/mattport/nerf/optimizers.py @@ -56,11 +56,10 @@ class Optimizers: def optimizer_step_all(self): """Run step for all optimizers.""" - # logging.info("optimizer_step_all") - for param_name, optimizer in self.optimizers.items(): + for _, optimizer in self.optimizers.items(): + # note that they key is the parameter name optimizer.step() def scheduler_step_all(self, step): """Run step for all schedulers.""" - # logging.info("optimizer_step_all") - pass + # TODO: this needs to be implemented diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 51349697a..eb9d6241d 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -5,6 +5,7 @@ import datetime import logging import os from time import time +from typing import Dict import torch import torch.distributed as dist @@ -32,6 +33,7 @@ class Trainer: self.local_rank = local_rank self.world_size = world_size # dataset variables + self.train_image_dataset = None self.train_dataset = None self.train_dataloader = None self.test_dataset = None @@ -66,7 +68,7 @@ class Trainer: self.train_dataset, batch_size=1, num_workers=self.config.dataloader.num_workers, - collate_fn=lambda batch_list: collate_batch_size_one(batch_list), + collate_fn=collate_batch_size_one, pin_memory=True, ) # TODO(ethan): implement the test dataset @@ -160,14 +162,10 @@ class Trainer: mssg += f"{k}: {v} | " print(mssg, end="\r") - def get_aggregated_loss(self, losses): - """_summary_ - - Args: - losses (_type_): _description_ - - Returns: - _type_: _description_ + @classmethod + def get_aggregated_loss(cls, losses: Dict[str, torch.tensor]): + """Returns the aggregated losses and the scalar for calling .backwards() on. + # TODO: move this out to another file/class/etc. """ loss_sum = 0.0 loss_dict = {} diff --git a/mattport/utils/misc.py b/mattport/utils/misc.py index 7d44ace6e..0ef07dcf5 100644 --- a/mattport/utils/misc.py +++ b/mattport/utils/misc.py @@ -1,3 +1,8 @@ +""" +Miscellaneous helper code. +""" + + class DotDict(dict): """ dot.notation access to dictionary attributes diff --git a/notebooks/visualize_dataset.ipynb b/notebooks/visualize_dataset.ipynb index 6f1dec38e..1012caefa 100644 --- a/notebooks/visualize_dataset.ipynb +++ b/notebooks/visualize_dataset.ipynb @@ -20,18 +20,19 @@ "import torch\n", "from torch.utils.data import DataLoader\n", "\n", - "from mattport.nerf.dataset.utils import get_dataset_inputs\n", "from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch\n", + "from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one\n", + "from mattport.nerf.dataset.utils import get_dataset_inputs\n", "from mattport.nerf.field_modules.ray_generator import RayGenerator\n", "from mattport.structures.rays import RayBundle\n", "from mattport.utils.io import get_absolute_path\n", + "from mattport.viewer.plotly import get_line_segments_from_lines\n", "\n", "from hydra import compose, initialize\n", "from omegaconf import open_dict\n", "import pprint\n", "from tqdm import tqdm\n", "\n", - "import goat\n", "import random\n", "from tqdm import tqdm" ] @@ -49,16 +50,23 @@ " \n", "cfg.dataset.data_directory = get_absolute_path(cfg.dataset.data_directory)\n", "dataset_inputs = get_dataset_inputs(**cfg.dataset)\n", - "dataset = ImageDataset(\n", - " image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor\n", + "image_dataset = ImageDataset(\n", + " image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor\n", + " )\n", + "iter_dataset = CollateIterDataset(\n", + " image_dataset,\n", + " collate_fn=lambda batch_list: collate_batch(\n", + " batch_list, cfg.dataloader.num_rays_per_batch, keep_full_image=True\n", + " ),\n", + " num_samples_to_collate=cfg.dataloader.num_images_to_sample_from,\n", + " num_times_to_repeat=cfg.dataloader.num_times_to_repeat_images,\n", ")\n", - "\n", "dataloader = DataLoader(\n", - " dataset,\n", - " batch_size=cfg.dataloader.num_images_to_sample_from,\n", - " collate_fn=lambda batch: collate_batch(batch, cfg.dataloader.num_rays_per_batch, keep_full_image=True),\n", + " iter_dataset,\n", + " batch_size=1,\n", " num_workers=cfg.dataloader.num_workers,\n", - " shuffle=True,\n", + " collate_fn=collate_batch_size_one,\n", + " pin_memory=True,\n", ")\n", "dataloader_iter = iter(dataloader)\n", "\n", @@ -67,7 +75,7 @@ "num_batches = 10\n", "for _ in tqdm(range(num_batches)):\n", " batch = next(dataloader_iter)\n", - " ray_bundle = ray_generator.forward(batch.indices)\n", + " ray_bundle = ray_generator.forward(batch[\"indices\"])\n", " break" ] }, @@ -78,7 +86,7 @@ "metadata": {}, "outputs": [], "source": [ - "media.show_image(dataset.get_image(10))" + "media.show_image(image_dataset.get_image(10))" ] }, { @@ -90,11 +98,11 @@ "source": [ "def show_batch(batch):\n", " # set the color of the sampled rays\n", - " c, y, x = [i.flatten() for i in torch.split(batch.indices, 1, dim=-1)]\n", - " batch.image[c,y,x] = 1.0\n", + " c, y, x = [i.flatten() for i in torch.split(batch[\"local_indices\"], 1, dim=-1)]\n", + " batch[\"image\"][c,y,x] = 1.0\n", "\n", - " # batch.image is num_images, h, w, 3\n", - " images = torch.split(batch.image, 1, dim=0)\n", + " # batch[\"image\"] is num_images, h, w, 3\n", + " images = torch.split(batch[\"image\"], 1, dim=0)\n", " image_list = [image[0] for image in images]\n", " image = torch.cat(image_list, dim=1) # cat along the width dimension\n", "\n", @@ -117,16 +125,6 @@ "batch = sample_and_show_batch()" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "f5f2e2e8-8d12-4e12-9b9f-e70a4f25bcee", - "metadata": {}, - "outputs": [], - "source": [ - "# batch.keys()" - ] - }, { "cell_type": "code", "execution_count": null, @@ -134,7 +132,7 @@ "metadata": {}, "outputs": [], "source": [ - "ray_indices = batch.indices\n", + "ray_indices = batch[\"indices\"]\n", "ray_bundle = ray_generator(ray_indices)" ] }, @@ -180,7 +178,7 @@ "metadata": {}, "outputs": [], "source": [ - "data += goat.plotly_utils.get_line_segments_from_lines(lines)" + "data += get_line_segments_from_lines(lines)" ] }, { @@ -240,6 +238,14 @@ "metadata": {}, "outputs": [], "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "722f5370-3b0f-4ad1-87a1-df479ae85b60", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/tests/nerf/dataloader/test_ray_generator.py b/tests/nerf/dataloader/test_ray_generator.py index a03d1c970..108a34a7f 100644 --- a/tests/nerf/dataloader/test_ray_generator.py +++ b/tests/nerf/dataloader/test_ray_generator.py @@ -4,10 +4,10 @@ Code to test ray generation. import imageio import torch -from hydra import compose, initialize from torch.utils.data import DataLoader from tqdm import tqdm +from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch from mattport.nerf.dataset.utils import get_dataset_inputs from mattport.nerf.field_modules.ray_generator import RayGenerator @@ -19,11 +19,11 @@ def visualize_batch(batch): """Returns an image from a batch.""" assert "image" in batch, "You must call collate_batch with keep_full_image=True for this to work." # set the color of the sampled rays - c, y, x = [i.flatten() for i in torch.split(batch.indices, 1, dim=-1)] - batch.image[c, y, x] = 1.0 + c, y, x = [i.flatten() for i in torch.split(batch["local_indices"], 1, dim=-1)] + batch["image"][c, y, x] = 1.0 - # batch.image is num_images, h, w, 3 - images = torch.split(batch.image, 1, dim=0) + # batch["image"] is num_images, h, w, 3 + images = torch.split(batch["image"], 1, dim=0) image_list = [image[0] for image in images] image = torch.cat(image_list, dim=1) # cat along the width dimension image = (image * 255.0).to(torch.uint8) @@ -33,23 +33,33 @@ def visualize_batch(batch): def test_dataloader(): """Testing for the dataloader from input dataset parameters to rays.""" - with initialize(config_path="configs"): - cfg = compose(config_name="blender_lego.yaml") + data_directory = "data/blender/lego" + dataset_type = "blender" + downscale_factor = 1 + num_images_to_sample_from = 8 + num_times_to_repeat_images = 40 + num_rays_per_batch = 1024 + num_workers = 0 - cfg.dataset.data_directory = get_absolute_path( - cfg.dataset.data_directory - ) # TODO(ethan): create a helper function for this - dataset_inputs = get_dataset_inputs(**cfg.dataset) - dataset = ImageDataset( + data_directory = get_absolute_path(data_directory) + dataset_inputs = get_dataset_inputs( + data_directory=data_directory, dataset_type=dataset_type, downscale_factor=downscale_factor + ) + image_dataset = ImageDataset( image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor ) - + iter_dataset = CollateIterDataset( + image_dataset, + collate_fn=lambda batch_list: collate_batch(batch_list, num_rays_per_batch, keep_full_image=True), + num_samples_to_collate=num_images_to_sample_from, + num_times_to_repeat=num_times_to_repeat_images, + ) dataloader = DataLoader( - dataset, - batch_size=cfg.dataloader.num_images_to_sample_from, - collate_fn=lambda batch: collate_batch(batch, cfg.dataloader.num_rays_per_batch, keep_full_image=True), - num_workers=cfg.dataloader.num_workers, - shuffle=True, + iter_dataset, + batch_size=1, + num_workers=num_workers, + collate_fn=collate_batch_size_one, + pin_memory=True, ) dataloader_iter = iter(dataloader) @@ -58,7 +68,7 @@ def test_dataloader(): num_batches = 10 for _ in tqdm(range(num_batches)): batch = next(dataloader_iter) - ray_bundle = ray_generator.forward(batch.indices) + ray_bundle = ray_generator.forward(batch["indices"]) # visualize the batch image = visualize_batch(batch)