From 3abc586427009072f7eeb975f26342c5fc7b20a6 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Mon, 1 Aug 2022 15:01:39 -0700 Subject: [PATCH] Remove some type ignores (#238) --- nerfactory/engine/trainer.py | 39 +++++++++-------- nerfactory/graphs/base.py | 85 +++++++++++++++++++++++------------- 2 files changed, 75 insertions(+), 49 deletions(-) diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index a21862942..0662f6379 100644 --- a/nerfactory/engine/trainer.py +++ b/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. diff --git a/nerfactory/graphs/base.py b/nerfactory/graphs/base.py index 955fd8e3f..e9512dd68 100644 --- a/nerfactory/graphs/base.py +++ b/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"""