Remove some type ignores (#238)

nerfactory/engine/trainer.py

@@ -18,7 +18,8 @@ Code to train model.
 import functools
 import logging
 import os
-from typing import Dict, List, Optional, Union
+import typing
+from typing import Dict, List, Optional
 
 import torch
 import torch.distributed as dist
@@ -69,7 +70,7 @@ class Trainer:
         self.dataloader_train: TrainDataloader
         self.dataloader_eval: EvalDataloader
         # model variables
-        self.graph: Union[Graph, DDP]
+        self.graph: Graph
         self.optimizers: Optimizers
         self.start_step = 0
         # logging variables
@@ -97,10 +98,10 @@ class Trainer:
         self._load_checkpoint()
 
         if self.world_size > 1:
-            self.graph = DDP(self.graph, device_ids=[self.local_rank])
+            self.graph = typing.cast(Graph, DDP(self.graph, device_ids=[self.local_rank]))
             dist.barrier(device_ids=[self.local_rank])
 
-        self.callbacks = self.graph.get_training_callbacks()  # type: ignore
+        self.callbacks = self.graph.get_training_callbacks()
 
     @classmethod
     def get_aggregated_loss(cls, loss_dict: Dict[str, torch.Tensor]) -> torch.Tensor:
@@ -135,7 +136,7 @@ class Trainer:
                 if step % self.config.trainer.steps_per_test == 0:
                     self.eval_with_dataloader(self.dataloader_eval, step=step)
                 self._write_out_storage(step)
-                self.visualizer_state.update_scene(step, self.graph)  # type: ignore
+                self.visualizer_state.update_scene(step, self.graph)
 
         self._write_out_storage(num_iterations)
 
@@ -162,7 +163,7 @@ class Trainer:
             loaded_state = torch.load(load_path, map_location="cpu")
             self.start_step = loaded_state["step"] + 1
             # load the checkpoints for graph and optimizer
-            self.graph.load_graph(loaded_state)  # type: ignore
+            self.graph.load_graph(loaded_state)
             self.optimizers.load_optimizers(loaded_state)
             self.grad_scaler.load_state_dict(loaded_state["scaler"])
             logging.info("done loading checkpoint from %s", load_path)
@@ -183,7 +184,7 @@ class Trainer:
         if hasattr(self.graph, "module"):
             model = self.graph.module.state_dict()  # type: ignore
         else:
-            model = self.graph.state_dict()  # type: ignore
+            model = self.graph.state_dict()
         torch.save(
             {
                 "step": step,
@@ -195,7 +196,9 @@ class Trainer:
         )
 
     @profiler.time_function
-    def train_iteration(self, ray_indices: TensorType["num_rays", 3], batch: dict, step: int) -> Dict[str, float]:
+    def train_iteration(
+        self, ray_indices: TensorType["num_rays", 3], batch: Dict[str, torch.Tensor], step: int
+    ) -> Dict[str, torch.Tensor]:
         """Run one iteration with a batch of inputs.
 
         Args:
@@ -208,8 +211,8 @@ class Trainer:
         """
         self.optimizers.zero_grad_all()
         with torch.autocast(device_type=ray_indices.device.type, enabled=self.mixed_precision):
-            _, loss_dict, metrics_dict = self.graph.forward(ray_indices=ray_indices, batch=batch)  # type: ignore
-            loss = sum(loss_dict.values())  # type: ignore
+            _, loss_dict, metrics_dict = self.graph.forward(ray_indices=ray_indices, batch=batch)
+            loss = typing.cast(torch.Tensor, sum(loss_dict.values()))
         self.grad_scaler.scale(loss).backward()  # type: ignore
         self.optimizers.optimizer_scaler_step_all(self.grad_scaler)
         self.grad_scaler.update()
@@ -219,9 +222,9 @@ class Trainer:
             callback.after_step(step)
 
         # Merging loss and metrics dict into a single output.
-        loss_dict["loss"] = loss  # type: ignore
-        loss_dict.update(metrics_dict)  # type: ignore
-        return loss_dict  # type: ignore
+        loss_dict["loss"] = loss
+        loss_dict.update(metrics_dict)
+        return loss_dict
 
     @profiler.time_function
     def test_image(self, camera_ray_bundle: RayBundle, batch: dict, step: Optional[int] = None) -> float:
@@ -235,14 +238,14 @@ class Trainer:
         Returns:
             float: PSNR
         """
-        self.graph.eval()  # type: ignore
+        self.graph.eval()
         if camera_ray_bundle.camera_indices is None:
             raise ValueError("camera_ray_bundle.camera_indices is None during testing")
         image_idx = int(camera_ray_bundle.camera_indices[0, 0])
-        outputs = self.graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle)  # type: ignore
-        psnr = self.graph.log_test_image_outputs(image_idx, step, batch, outputs)  # type: ignore
-        self.graph.train()  # type: ignore
-        return psnr  # type: ignore
+        outputs = self.graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
+        psnr = self.graph.log_test_image_outputs(image_idx, step, batch, outputs)
+        self.graph.train()
+        return psnr
 
     def eval_with_dataloader(self, dataloader: EvalDataloader, step: Optional[int] = None) -> None:
         """Run evaluation with a given dataloader.

nerfactory/graphs/base.py

@@ -17,7 +17,7 @@ The Graph module contains all trainable parameters.
 """
 from abc import abstractmethod
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union, overload
 
 import torch
 from omegaconf import DictConfig
@@ -39,27 +39,7 @@ from nerfactory.utils.misc import (
 )
 
 
-class AbstractGraph(nn.Module):
-    """Highest level graph class. Somewhat useful to lift code up and out of the way."""
-
-    def __init__(self) -> None:
-        super().__init__()
-        # to keep track of which device the nn.Module is on
-        self.device_indicator_param = nn.Parameter(torch.empty(0))
-
-    @property
-    def device(self):
-        """Returns the device that the graph is on."""
-        return self.device_indicator_param.device
-
-    @abstractmethod
-    def forward(
-        self, ray_indices: TensorType["num_rays", 3], batch: Union[str, Optional[Dict[str, torch.Tensor]]] = None
-    ):
-        """Process starting with ray indices. Turns them into rays, then performs volume rendering."""
-
-
-class Graph(AbstractGraph):
+class Graph(nn.Module):
     """Graph class
     Where everything (Fields, Optimizers, Samplers, Visualization, etc) is linked together. This should be
     subclassed for custom NeRF model.
@@ -109,6 +89,13 @@ class Graph(AbstractGraph):
         # variable for visualizer to fetch TODO(figure out if there is cleaner way to do this)
         self.vis_outputs = None
         self.default_output_name = None
+        # to keep track of which device the nn.Module is on
+        self.device_indicator_param = nn.Parameter(torch.empty(0))
+
+    @property
+    def device(self):
+        """Returns the device that the graph is on."""
+        return self.device_indicator_param.device
 
     def get_training_callbacks(self) -> List[Callback]:  # pylint:disable=no-self-use
         """Returns a list of callbacks that run functions at the specified training iterations."""
@@ -159,9 +146,21 @@ class Graph(AbstractGraph):
                 v = torch.tile(v, (1, 1, 3))
             outputs[k] = v
 
+    @overload
+    def forward_after_ray_generator(self, ray_bundle: RayBundle, batch: None = None) -> Dict[str, torch.Tensor]:
+        ...
+
+    @overload
     def forward_after_ray_generator(
-        self, ray_bundle: RayBundle, batch: Optional[Union[str, Dict[str, torch.Tensor]]] = None
-    ):
+        self, ray_bundle: RayBundle, batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        ...
+
+    def forward_after_ray_generator(
+        self, ray_bundle: RayBundle, batch: Optional[Dict[str, torch.Tensor]] = None
+    ) -> Union[
+        Dict[str, torch.Tensor], Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, torch.Tensor]]
+    ]:
         """Run forward starting with a ray bundle."""
         if self.collider is not None:
             intersected_ray_bundle = self.collider(ray_bundle)
@@ -193,12 +192,25 @@ class Graph(AbstractGraph):
                 loss_dict[loss_name] *= self.loss_coefficients[loss_name]
         return outputs, loss_dict, metrics_dict
 
+    @overload
+    def forward(self, ray_indices: TensorType["num_rays", 3], batch: None = None) -> Dict[str, torch.Tensor]:
+        ...
+
+    @overload
     def forward(
-        self, ray_indices: TensorType["num_rays", 3], batch: Optional[Union[str, Dict[str, torch.Tensor]]] = None
-    ):
-        """Run the forward starting with ray indices."""
+        self, ray_indices: TensorType["num_rays", 3], batch: Dict[str, torch.Tensor]
+    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        ...
+
+    def forward(
+        self, ray_indices: TensorType["num_rays", 3], batch: Optional[Dict[str, torch.Tensor]] = None
+    ) -> Union[
+        Dict[str, torch.Tensor], Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, torch.Tensor]]
+    ]:
+        """Process starting with ray indices. Turns them into rays, then performs volume rendering."""
         ray_bundle = self.ray_generator.forward(ray_indices)
-        return self.forward_after_ray_generator(ray_bundle, batch=batch)
+        x = self.forward_after_ray_generator(ray_bundle, batch=batch)
+        return x
 
     def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]:
         """Compute and returns metrics."""
@@ -222,7 +234,7 @@ class Graph(AbstractGraph):
             start_idx = i
             end_idx = i + camera_ray_bundle.num_rays_per_chunk
             ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx)
-            outputs = self.forward_after_ray_generator(ray_bundle)
+            outputs = self.forward_after_ray_generator(ray_bundle=ray_bundle)
             for output_name, output in outputs.items():  # type: ignore
                 outputs_lists[output_name].append(output)
         for output_name, outputs_list in outputs_lists.items():
@@ -235,8 +247,19 @@ class Graph(AbstractGraph):
         return self.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
 
     @abstractmethod
-    def log_test_image_outputs(self, image_idx, step, batch, outputs):
-        """Writes the test image outputs."""
+    def log_test_image_outputs(self, image_idx, step, batch, outputs) -> float:
+        """Writes the test image outputs.
+        TODO: This shouldn't return a loss
+
+        Args:
+            image_idx: Index of the image.
+            step: Current step.
+            batch: Batch of data.
+            outputs: Outputs of the graph.
+
+        Returns:
+            The psnr.
+        """
 
     def load_graph(self, loaded_state: Dict[str, Any]) -> None:
         """Load the checkpoint from the given path"""