From 7e5a5bf524d6e44b2a6190cf6ae323424a36e832 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Thu, 21 Apr 2022 21:57:36 +0000 Subject: [PATCH] code runs again, working on integrating test set and optimizers --- mattport/nerf/dataset/image_dataset.py | 2 +- mattport/nerf/dataset/utils.py | 12 ------------ mattport/nerf/field/nerf.py | 6 ++++-- mattport/nerf/graph/nerf.py | 11 +++++++++-- mattport/nerf/optimizers.py | 20 +++++++++++++++++++- mattport/nerf/trainer.py | 20 +++++++++++++++++--- mattport/utils/misc.py | 10 ++++++++++ 7 files changed, 60 insertions(+), 21 deletions(-) create mode 100644 mattport/utils/misc.py diff --git a/mattport/nerf/dataset/image_dataset.py b/mattport/nerf/dataset/image_dataset.py index 2ef3ebfeb..a2f5a6581 100644 --- a/mattport/nerf/dataset/image_dataset.py +++ b/mattport/nerf/dataset/image_dataset.py @@ -9,7 +9,7 @@ import torch from PIL import Image from torch.utils.data import default_collate -from mattport.nerf.dataset.utils import DotDict +from mattport.utils.misc import DotDict class ImageDataset(torch.utils.data.Dataset): diff --git a/mattport/nerf/dataset/utils.py b/mattport/nerf/dataset/utils.py index f7c02ac9f..afeff467c 100644 --- a/mattport/nerf/dataset/utils.py +++ b/mattport/nerf/dataset/utils.py @@ -13,18 +13,6 @@ import torch from mattport.utils.io import get_absolute_path, load_from_json -class DotDict(dict): - """ - dot.notation access to dictionary attributes - """ - - def __getattr__(self, attr): - return self[attr] - - __setattr__ = dict.__setitem__ - __delattr__ = dict.__delitem__ - - @dataclass class DatasetInputs: """Dataset inputs for the image dataset and the ray generator.""" diff --git a/mattport/nerf/field/nerf.py b/mattport/nerf/field/nerf.py index 52e9c7880..ba519d6f3 100644 --- a/mattport/nerf/field/nerf.py +++ b/mattport/nerf/field/nerf.py @@ -52,5 +52,7 @@ class NeRFField(nn.Module): field_rgb_output = self.field_output_rgb(rgb_mlp_out) field_density_out = self.field_output_density(base_mlp_out) - field_outputs = field_rgb_output | field_density_out - return field_outputs + field_outputs = {} + field_outputs.update(field_rgb_output) + field_outputs.update(field_density_out) + return field_outputs \ No newline at end of file diff --git a/mattport/nerf/graph/nerf.py b/mattport/nerf/graph/nerf.py index 20af5aa7b..8289b18b1 100644 --- a/mattport/nerf/graph/nerf.py +++ b/mattport/nerf/graph/nerf.py @@ -13,6 +13,7 @@ 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.field_modules.field_heads import FieldHeadNames class NeRFGraph(Graph): @@ -61,14 +62,20 @@ class NeRFGraph(Graph): # coarse network: uniform_ray_samples = self.sampler_uniform(ray_bundle) # RaySamples coarse_field_outputs = self.field_coarse(uniform_ray_samples) # FieldOutputs + coarse_renderer_outputs = self.renderer_rgb( - field_outputs=coarse_field_outputs.rgb, deltas=uniform_ray_samples.deltas + rgb=coarse_field_outputs[FieldHeadNames.RGB], + density=coarse_field_outputs[FieldHeadNames.DENSITY], + deltas=uniform_ray_samples.deltas, ) # RendererOutputs # fine network: pdf_ray_samples = self.sampler_pdf(ray_bundle, uniform_ray_samples, coarse_field_outputs) # RaySamples fine_field_outputs = self.field_fine(pdf_ray_samples) # FieldOutputs + fine_renderer_outputs = self.renderer_rgb( - field_outputs=fine_field_outputs, deltas=pdf_ray_samples.deltas + rgb=fine_field_outputs[FieldHeadNames.RGB], + density=fine_field_outputs[FieldHeadNames.DENSITY], + deltas=pdf_ray_samples.deltas, ) # RendererOutputs # outputs: outputs = {"rgb_coarse": coarse_renderer_outputs.rgb, "rgb_fine": fine_renderer_outputs.rgb} diff --git a/mattport/nerf/optimizers.py b/mattport/nerf/optimizers.py index 6e59836ca..6490dc6a0 100644 --- a/mattport/nerf/optimizers.py +++ b/mattport/nerf/optimizers.py @@ -2,6 +2,8 @@ Optimizers class. """ +import logging + from typing import Dict, List from omegaconf import DictConfig import torch @@ -47,4 +49,20 @@ class Optimizers: param_name (str): name of scheduler to step forward """ if self.config.param_name.scheduler: - self.schedulers[param_name].step() \ No newline at end of file + self.schedulers[param_name].step() + + def zero_grad_all(self): + """zero the gradients for all optimizer parameters""" + for param_name, optimizer in self.optimizers.items(): + # print(param_name) + pass + + def optimizer_step_all(self): + """Run step for all optimizers.""" + # logging.info("optimizer_step_all") + pass + + def scheduler_step_all(self): + """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 0e9eebd1f..3fb67b594 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -74,6 +74,7 @@ class Trainer: def setup_optimizers(self): """_summary_""" + # TODO(ethan): handle different world sizes self.optimizers = Optimizers(self.config.param_groups, self.graph.get_param_groups()) def load_checkpoint(self, load_config: DictConfig) -> int: @@ -116,6 +117,13 @@ class Trainer: ckpt_path, ) + def get_aggregated_loss(self, losses): + aggregated_loss = 0.0 + for loss_name in losses.keys(): + # TODO(ethan): add loss weightings here from a config + aggregated_loss += losses[loss_name] + return aggregated_loss + def train(self) -> None: """_summary_""" num_iterations = 1000 @@ -123,7 +131,7 @@ class Trainer: 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.save_step == 0: + if 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") @@ -133,11 +141,17 @@ class Trainer: 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}") - losses = self.graph.module.get_losses(batch, graph_outputs) + # print(batch) + losses = self.graph.get_losses(batch, graph_outputs) # or self.graph.module + loss = self.get_aggregated_loss(losses) + self.optimizers.zero_grad_all() + loss.backward() + self.optimizers.scheduler_step_all() + self.optimizers.optimizer_step_all() return losses def test(self, image_idx): """Test a specific image.""" with torch.no_grad(): ray_indices = None - \ No newline at end of file + pass diff --git a/mattport/utils/misc.py b/mattport/utils/misc.py new file mode 100644 index 000000000..d6fb5ecfe --- /dev/null +++ b/mattport/utils/misc.py @@ -0,0 +1,10 @@ +class DotDict(dict): + """ + dot.notation access to dictionary attributes + """ + + def __getattr__(self, attr): + return self[attr] + + __setattr__ = dict.__setitem__ + __delattr__ = dict.__delitem__ \ No newline at end of file