check DDP running correctly

configs/config.yaml

@@ -17,7 +17,7 @@ optimizer:
 
 # machine options
 machine_config:
-  num_gpus: 1
+  num_gpus: 4
   num_machines: 1
   machine_rank: 0
   dist_url: auto

mattport/nerf/graph/base.py

@@ -31,35 +31,13 @@ class Node:
         return hash(self.name)
 
 
-class Graph(nn.ModuleDict):
+class Graph(nn.Module):
     """_summary_"""
 
     def __init__(self, intrinsics=None, camera_to_world=None) -> None:
         super().__init__()
         self.intrinsics = intrinsics
         self.camera_to_world = camera_to_world
-
-        self.modules = None  # NOTE(ethan): I'm turning off these features for now
-        if self.modules is not None:
-            # create graph and get ordering
-            self.roots = self.construct_graph()
-            self.module_order = self.get_module_order()
-
-            # initialize graph with known input dimensions; set default in_dim to 0
-            for module_name, module_dict in self.modules.items():
-                module = getattr(importlib.import_module("mattport.nerf.field_modules"), module_dict.class_name)
-                if not module_dict.meta_data.in_dim:
-                    module_dict.meta_data.in_dim = 0
-                self[module_name] = module(**module_dict.meta_data)
-
-            # calculate input dimensions based on module dependencies
-            for root in self.roots:
-                self.get_in_dim(root)
-
-            # instantiate torch.nn members of network
-            for _, module in self.items():
-                module.build_nn_modules()
-
         self.populate_modules()  # populate the modules
 
     @abstractmethod

mattport/nerf/trainer.py

@@ -17,6 +17,7 @@ from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch
 from mattport.nerf.dataset.utils import get_dataset_inputs
 from mattport.utils.decorators import check_main_thread
 
+logging.getLogger('PIL').setLevel(logging.WARNING)
 
 class Trainer:
     """Training class"""
@@ -129,5 +130,5 @@ 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.get_losses(batch, graph_outputs)
-        logging.info(losses)
+        losses = self.graph.module.get_losses(batch, graph_outputs)
+        # logging.info(losses)