callback for optimizer (#99)

mattport/nerf/occupancy_grid.py

@@ -103,13 +103,11 @@ class OccupancyGrid(nn.Module):
         self.mean_density = torch.mean(self.occupancy_grid[valid_mask]).item()
 
     @torch.no_grad()
-    def get_densities(self, xyzs: TensorType[..., 3], update_iter_count: bool = False) -> TensorType[..., 1]:
+    def get_densities(self, xyzs: TensorType[..., 3]) -> TensorType[..., 1]:
         """Trilinear interpolation to get the density values.
 
         Args:
             xyzs (TensorType[..., 3]): 3D querry coordinate
-            update_iter_count (bool, optional): Update current iteration. Used for knowing when to update grid.
-                Defaults to False.
 
         Returns:
             TensorType[..., 1]: Density values
@@ -128,14 +126,6 @@ class OccupancyGrid(nn.Module):
             padding_mode="zeros",
         )
         densities = densities.view(*xyzs_shape[:-1], 1)
-
-        # TODO we should be tracking the iteration another way
-        if self.training and update_iter_count:
-            self.iteration_count += 1
-            if self.iteration_count % self.update_every_num_iters == 0:
-                if self.density_fn is not None:
-                    self.update_occupancy_grid(self.density_fn)
-
         return densities
 
     def forward(self, x):

mattport/nerf/sampler.py

@@ -32,7 +32,7 @@ class Sampler(nn.Module):
         """Generate ray samples with optional occupancy filtering"""
         ray_samples = self.generate_ray_samples(*args, **kwargs)
         if self.occupancy_field is not None:
-            densities = self.occupancy_field.get_densities(ray_samples.positions, update_iter_count=True)
+            densities = self.occupancy_field.get_densities(ray_samples.positions)
             weights = ray_samples.get_weights(densities)
 
             valid_mask = weights >= self.weight_threshold

mattport/nerf/trainer.py

@@ -6,7 +6,7 @@ import logging
 import os
 from pydoc import locate
 from time import time
-from typing import Dict
+from typing import Callable, Dict, List
 
 import torch
 import torch.distributed as dist
@@ -19,6 +19,7 @@ from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch
 from mattport.nerf.dataset.utils import DatasetInputs, get_dataset_inputs
 from mattport.nerf.optimizers import Optimizers
 from mattport.utils import profiler, writer
+from mattport.utils.callbacks import update_occupancy
 from mattport.utils.decorators import check_main_thread
 
 logging.getLogger("PIL").setLevel(logging.WARNING)
@@ -210,7 +211,7 @@ class Trainer:
             )
 
             iter_start = time()
-            loss_dict = self.train_iteration(batch, step)
+            loss_dict = self.train_iteration(batch, step, _callback=[update_occupancy])
             writer.put_time(
                 name=writer.EventName.RAYS_PER_SEC,
                 start_time=iter_start,
@@ -254,7 +255,7 @@ class Trainer:
             writer.write_out_storage()
 
     @profiler.time_function
-    def train_iteration(self, batch: dict, step: int):
+    def train_iteration(self, batch: dict, step: int, _callback: List[Callable] = None):
         """Run one iteration with a batch of inputs."""
         # move batch to correct device
         ray_indices = batch["indices"].to(self.device)
@@ -264,6 +265,9 @@ class Trainer:
         loss.backward()
         self.optimizers.optimizer_step_all()
         self.optimizers.scheduler_step_all(step)
+        if _callback:
+            for _func in _callback:
+                _func(self.graph)
         return loss_dict
 
     @profiler.time_function

mattport/utils/callbacks.py

@@ -0,0 +1,21 @@
+"""
+Callback functions for training iterations
+"""
+
+from mattport.nerf.graph.base import Graph
+from mattport.nerf.occupancy_grid import OccupancyGrid
+
+
+def update_occupancy(graph: Graph):
+    """Update the occupancy grid depending on iteration count after every train step"""
+    if not hasattr(graph, "occupancy_pointers"):
+        occupancy_grids = []
+        for _, module in graph.__dict__["_modules"].items():
+            if isinstance(module, OccupancyGrid):
+                occupancy_grids.append(module)
+        graph.occupancy_pointers = occupancy_grids
+    for grid in graph.occupancy_pointers:
+        grid.iteration_count += 1
+        if grid.iteration_count % grid.update_every_num_iters == 0:
+            if grid.density_fn is not None:
+                grid.update_occupancy_grid(grid.density_fn)

mattport/utils/writer.py

@@ -320,15 +320,18 @@ class LocalWriter(Writer):
 
     def _update_header(self, latest_map, new_key):
         """helper to handle the printing of the header labels"""
-        if (not self.max_log_size and len(self.past_mssgs) == 0) or (
-            self.max_log_size and (len(self.past_mssgs) == 0 or new_key)
-        ):
+        full_log_cond = not self.max_log_size and GLOBAL_BUFFER["step"] <= GLOBAL_BUFFER["steps_per_log"]
+        capped_log_cond = self.max_log_size and (len(self.past_mssgs) <= 2 or new_key)
+        if full_log_cond or capped_log_cond:
             mssg = f"{'Step (% Done)':<20}"
             for name, _ in latest_map.items():
                 if name in self.stats_to_track:
                     mssg += f"{name:<20} "
             self.past_mssgs[0] = mssg
             self.past_mssgs[1] = "-" * len(mssg)
+            if full_log_cond:
+                print(mssg)
+                print("-" * len(mssg))
 
     def _print_stats(self, latest_map, padding=" "):
         """helper to print out the stats in a readable format"""