From cf7034f0d1ff89c8df591c145bcac5450885ae67 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Fri, 22 Apr 2022 03:45:06 +0000 Subject: [PATCH] added test set (using training data for now) --- .vscode/launch.json | 6 +- README.md | 5 +- configs/{config.yaml => default.yaml} | 21 ++++-- mattport/nerf/dataset/image_dataset.py | 2 +- mattport/nerf/graph/nerf.py | 4 +- mattport/nerf/metrics.py | 13 ++++ mattport/nerf/optimizers.py | 12 ++-- mattport/nerf/trainer.py | 71 +++++++++++++++------ mattport/utils/writer.py | 33 +++++----- scripts/{run_train_nerf.py => run_train.py} | 2 +- 10 files changed, 112 insertions(+), 57 deletions(-) rename configs/{config.yaml => default.yaml} (70%) create mode 100644 mattport/nerf/metrics.py rename scripts/{run_train_nerf.py => run_train.py} (98%) diff --git a/.vscode/launch.json b/.vscode/launch.json index 13aa61c4c..0cae34962 100644 --- a/.vscode/launch.json +++ b/.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" + ] } ] } \ No newline at end of file diff --git a/README.md b/README.md index fe3586272..f8c1f7179 100644 --- a/README.md +++ b/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 diff --git a/configs/config.yaml b/configs/default.yaml similarity index 70% rename from configs/config.yaml rename to configs/default.yaml index 397a54279..331f6880c 100644 --- a/configs/config.yaml +++ b/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 diff --git a/mattport/nerf/dataset/image_dataset.py b/mattport/nerf/dataset/image_dataset.py index a2f5a6581..7acafa66a 100644 --- a/mattport/nerf/dataset/image_dataset.py +++ b/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: diff --git a/mattport/nerf/graph/nerf.py b/mattport/nerf/graph/nerf.py index 8289b18b1..a565e5dc6 100644 --- a/mattport/nerf/graph/nerf.py +++ b/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]): diff --git a/mattport/nerf/metrics.py b/mattport/nerf/metrics.py new file mode 100644 index 000000000..230b66327 --- /dev/null +++ b/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) diff --git a/mattport/nerf/optimizers.py b/mattport/nerf/optimizers.py index 6490dc6a0..0db806fd5 100644 --- a/mattport/nerf/optimizers.py +++ b/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 \ No newline at end of file diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 3fb67b594..5f7926bad 100644 --- a/mattport/nerf/trainer.py +++ b/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") diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index 3b7dc7c79..d1e8cd520 100644 --- a/mattport/utils/writer.py +++ b/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 diff --git a/scripts/run_train_nerf.py b/scripts/run_train.py similarity index 98% rename from scripts/run_train_nerf.py rename to scripts/run_train.py index abe58c7ae..a93f51c08 100644 --- a/scripts/run_train_nerf.py +++ b/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(