From af74b4e95ceab826debdeabea9255b5c9aea43d0 Mon Sep 17 00:00:00 2001
From: evonneng <evonne_ng@berkeley.edu>
Date: Thu, 21 Apr 2022 17:06:29 +0000
Subject: [PATCH] check DDP running correctly

---
 configs/config.yaml         |  2 +-
 mattport/nerf/graph/base.py | 24 +-----------------------
 mattport/nerf/trainer.py    |  5 +++--
 3 files changed, 5 insertions(+), 26 deletions(-)

diff --git a/configs/config.yaml b/configs/config.yaml
index 8693a0ccb..85f57247b 100644
--- a/configs/config.yaml
+++ b/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
diff --git a/mattport/nerf/graph/base.py b/mattport/nerf/graph/base.py
index 89213c164..750dd3072 100644
--- a/mattport/nerf/graph/base.py
+++ b/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
diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index 129c81fdf..74e7d01f5 100644
--- a/mattport/nerf/trainer.py
+++ b/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)