code runs again, working on integrating test set and optimizers

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):

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."""

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

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}

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()
+            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

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
-            
+            pass

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__