From 7dc9d445b2e047f707d1ea6d0ed71f7bc4fcc6ab Mon Sep 17 00:00:00 2001
From: evonneng <evonne_ng@berkeley.edu>
Date: Fri, 15 Apr 2022 01:04:17 +0000
Subject: [PATCH 1/6] added basic distributed logic for trainer

---
 mattport/nerf/trainer.py                 | 58 ++++++++++++++++++++++--
 scripts/run_train_nerf.py                |  6 +--
 tests/nerf/configs/test_basic_graph.yaml |  1 +
 tests/nerf/test_trainer.py               | 15 ++++++
 4 files changed, 71 insertions(+), 9 deletions(-)
 create mode 100644 tests/nerf/test_trainer.py

diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index b827960a4..4bf367231 100644
--- a/mattport/nerf/trainer.py
+++ b/mattport/nerf/trainer.py
@@ -1,6 +1,16 @@
 """
 Code to train model.
 """
+import os
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from omegaconf import DictConfig
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data.distributed import DistributedSampler
+
 from mattport.nerf.dataset.build import build_dataset
 from mattport.nerf.optimizer import Optimizer
 
@@ -44,14 +54,52 @@ class Trainer:
         """
         raise NotImplementedError
 
-    def train(self):
+    def train(self, gpu: int) -> None:
+        """_summary_"""
+        self.setup_ddp(gpu, self.config.world_size)
+
+        # TODO(): implement training pipeline. below I add skeleton code for compatibility with DDP
+        # self.setup_graph()
+        # self.graph = DDP(self.graph, device_ids=[gpu])
+        # self.setup_dataset()
+        # train_sampler = DistributedSampler(self.dataset, rank=gpu, num_replicas=self.config.world_size, shuffle=True)
+        # TODO(): create dataloader that takes as input the train_sampler
+        # self.setup_optimizer()
+
+        # TODO(): calls to train_epoch and train_iteration loop
+
+        dist.destroy_process_group()
+        raise NotImplementedError
+
+    def train_epoch(self, gpu: int):
         """_summary_"""
         raise NotImplementedError
 
-    def train_epoch(self):
+    def train_iteration(self, gpu: int):
         """_summary_"""
+        # TODO(): save checkpoints, and do logging only on rank 0 device only
+        # if gpu == 0:
         raise NotImplementedError
 
-    def train_iteration(self):
-        """_summary_"""
-        raise NotImplementedError
+    def setup_ddp(self, gpu: int, world_size: int) -> None:
+        """Initialization of processes for distributed training.
+
+        Args:
+            gpu (int): process rank number
+            world_size (int): total number of gpus available for training
+        """
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = "12356"
+        dist.init_process_group("nccl", init_method="env://", rank=gpu, world_size=world_size)
+        torch.cuda.set_device(gpu)
+
+        torch.manual_seed(0)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        np.random.seed(0)
+
+    def run_multiprocess_train(self) -> None:
+        """function to spawn multiple processes for training."""
+        if self.config.world_size < 1:
+            self.config.world_size = torch.cuda.device_count()
+        mp.spawn(self.train, nprocs=self.config.world_size)
diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py
index 2ad9a5f97..791020ac1 100644
--- a/scripts/run_train_nerf.py
+++ b/scripts/run_train_nerf.py
@@ -2,12 +2,12 @@
 run_train_nerf.py
 """
 import logging
+
 import hydra
 from omegaconf import DictConfig
 
 from mattport.nerf.trainer import Trainer
 
-
 logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG)
 
 
@@ -15,9 +15,7 @@ logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=loggin
 def main(config: DictConfig):
     """Main function."""
     trainer = Trainer(config)
-    trainer.setup_dataset()
-    trainer.setup_graph()
-    trainer.setup_optimizer()
+    trainer.run_multiprocess_train()
 
 
 if __name__ == "__main__":
diff --git a/tests/nerf/configs/test_basic_graph.yaml b/tests/nerf/configs/test_basic_graph.yaml
index b3e23bd79..77939a274 100644
--- a/tests/nerf/configs/test_basic_graph.yaml
+++ b/tests/nerf/configs/test_basic_graph.yaml
@@ -1,3 +1,4 @@
+world_size: 0
 network:
   mlp_0:
     class_name: MLP
diff --git a/tests/nerf/test_trainer.py b/tests/nerf/test_trainer.py
new file mode 100644
index 000000000..0cc9ff4dd
--- /dev/null
+++ b/tests/nerf/test_trainer.py
@@ -0,0 +1,15 @@
+from tty import CFLAG
+import pytest
+import runpy
+
+from mattport.nerf.trainer import Trainer
+from hydra import compose, initialize
+
+CONFIG_DIR = "configs"
+
+def test_multiprocess_train():
+    """Test for trainer to validate distributed data parallel working."""
+    with initialize(config_path=CONFIG_DIR):
+        cfg = compose(config_name='test_basic_graph')
+    trainer = Trainer(cfg)
+    trainer.run_multiprocess_train()
\ No newline at end of file

From ac3e7515ef416b719f27534b6c657bd62c833fcc Mon Sep 17 00:00:00 2001
From: evonneng <evonne_ng@berkeley.edu>
Date: Fri, 15 Apr 2022 21:33:50 +0000
Subject: [PATCH 2/6] done adding main logic in run_train_nerf to initialize
 processes with more robust handling of DDP

---
 mattport/nerf/trainer.py                 |  50 +------
 scripts/run_train_nerf.py                | 173 ++++++++++++++++++++++-
 tests/nerf/configs/test_basic_graph.yaml |   6 +-
 3 files changed, 182 insertions(+), 47 deletions(-)

diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index 4bf367231..549934573 100644
--- a/mattport/nerf/trainer.py
+++ b/mattport/nerf/trainer.py
@@ -8,8 +8,6 @@ import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 from omegaconf import DictConfig
-from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.utils.data.distributed import DistributedSampler
 
 from mattport.nerf.dataset.build import build_dataset
 from mattport.nerf.optimizer import Optimizer
@@ -54,52 +52,16 @@ class Trainer:
         """
         raise NotImplementedError
 
-    def train(self, gpu: int) -> None:
-        """_summary_"""
-        self.setup_ddp(gpu, self.config.world_size)
-
-        # TODO(): implement training pipeline. below I add skeleton code for compatibility with DDP
-        # self.setup_graph()
-        # self.graph = DDP(self.graph, device_ids=[gpu])
-        # self.setup_dataset()
-        # train_sampler = DistributedSampler(self.dataset, rank=gpu, num_replicas=self.config.world_size, shuffle=True)
-        # TODO(): create dataloader that takes as input the train_sampler
-        # self.setup_optimizer()
-
-        # TODO(): calls to train_epoch and train_iteration loop
-
-        dist.destroy_process_group()
-        raise NotImplementedError
-
-    def train_epoch(self, gpu: int):
+    def train(self) -> None:
         """_summary_"""
         raise NotImplementedError
 
-    def train_iteration(self, gpu: int):
+    def train_epoch(self):
+        """_summary_"""
+        raise NotImplementedError
+
+    def train_iteration(self):
         """_summary_"""
         # TODO(): save checkpoints, and do logging only on rank 0 device only
         # if gpu == 0:
         raise NotImplementedError
-
-    def setup_ddp(self, gpu: int, world_size: int) -> None:
-        """Initialization of processes for distributed training.
-
-        Args:
-            gpu (int): process rank number
-            world_size (int): total number of gpus available for training
-        """
-        os.environ["MASTER_ADDR"] = "localhost"
-        os.environ["MASTER_PORT"] = "12356"
-        dist.init_process_group("nccl", init_method="env://", rank=gpu, world_size=world_size)
-        torch.cuda.set_device(gpu)
-
-        torch.manual_seed(0)
-        torch.backends.cudnn.deterministic = True
-        torch.backends.cudnn.benchmark = False
-        np.random.seed(0)
-
-    def run_multiprocess_train(self) -> None:
-        """function to spawn multiple processes for training."""
-        if self.config.world_size < 1:
-            self.config.world_size = torch.cuda.device_count()
-        mp.spawn(self.train, nprocs=self.config.world_size)
diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py
index 791020ac1..90a1ab246 100644
--- a/scripts/run_train_nerf.py
+++ b/scripts/run_train_nerf.py
@@ -1,21 +1,190 @@
 """
 run_train_nerf.py
 """
+import datetime
 import logging
+import random
+import socket
+from datetime import timedelta
+from typing import Any, Callable
 
 import hydra
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
 from omegaconf import DictConfig
 
 from mattport.nerf.trainer import Trainer
 
 logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG)
 
+DEFAULT_TIMEOUT = timedelta(minutes=30)
+
+
+def _find_free_port() -> str:
+    """Find free port on open socket"""
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.bind(("", 0))
+    port = sock.getsockname()[1]
+    sock.close()
+    return port
+
+
+def _set_random_seed(seed) -> None:
+    """Set randomness seed in torch and numpy"""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def _synchronize(world_size, local_rank):
+    """Wait for all processes on node finish"""
+    if world_size > 1:
+        dist.barrier(device_ids=[local_rank])
+
+
+def _distributed_worker(
+    local_rank: int,
+    main_func: Callable,
+    world_size: int,
+    num_gpus_per_machine: int,
+    machine_rank: int,
+    dist_url: str,
+    config: DictConfig,
+    timeout: datetime = DEFAULT_TIMEOUT,
+) -> Any:
+    """Spawned distributed worker that handles the initialization of process group and handles the
+       training process on multiple processes
+
+    Args:
+        local_rank (int): current rank of process
+        main_func (Callable): function that will be called by the distributed workers
+        world_size (int): total number of gpus available
+        num_gpus_per_machine (int): number of GPUs per machine
+        machine_rank (int): rank of this machine
+        dist_url (str): url to connect to for distributed jobs, including protocol
+                        e.g. "tcp://127.0.0.1:8686".
+                        Can be set to "auto" to automatically select a free port on localhost
+        config (DictConfig): config file specifying trainng regimen
+        timeout (datetime, optional): timeout of the distributed workers Defaults to DEFAULT_TIMEOUT.
+
+    Raises:
+        e: Exception in initializing the process group
+
+    Returns:
+        Any: TODO(): determine the return type
+    """
+    assert torch.cuda.is_available(), "cuda is not available. Please check your installation."
+    global_rank = machine_rank * num_gpus_per_machine + local_rank
+    try:
+        dist.init_process_group(
+            backend="NCCL",
+            init_method=dist_url,
+            world_size=world_size,
+            rank=global_rank,
+            timeout=timeout,
+        )
+    except Exception as e:
+        logger = logging.getLogger(__name__)
+        logger.error("Process group URL: %s", dist_url)
+        raise e
+
+    assert num_gpus_per_machine <= torch.cuda.device_count()
+    torch.cuda.set_device(local_rank)
+    _set_random_seed(1234 + local_rank)
+    _synchronize(world_size, local_rank)
+
+    output = main_func(local_rank, world_size, config)
+    _synchronize(world_size, local_rank)
+    dist.destroy_process_group()
+    return output
+
+
+def _train(local_rank: int, world_size: int, config: DictConfig) -> Any:
+    """Main training function that sets up and runs the trainer per process
+
+    Args:
+        local_rank (int): _description_
+        world_size (int): _description_
+        config (DictConfig): _description_
+
+    Returns:
+        Any: TODO(): determine the return type
+    """
+    # trainer = Trainer(local_rank, world_size, config)
+    # trainer.train()
+    return 0
+
+
+def launch(
+    main_func: Callable,
+    num_gpus_per_machine: int,
+    num_machines: int = 1,
+    machine_rank: int = 0,
+    dist_url: str = "auto",
+    config: DictConfig = None,
+    timeout: datetime = DEFAULT_TIMEOUT,
+) -> None:
+    """Function that spawns muliple processes to call on main_func
+
+    Args:
+        main_func (Callable): function that will be called by the distributed workers
+        num_gpus_per_machine (int): number of GPUs per machine
+        num_machines (int, optional): total number of machines
+        machine_rank (int, optional): rank of this machine. Defaults to 0.
+        dist_url (str, optional): url to connect to for distributed jobs. Defaults to "auto".
+        config (DictConfig, optional): config file specifying trainng regimen Defaults to None.
+        timeout (datetime, optional): timeout of the distributed workers Defaults to DEFAULT_TIMEOUT.
+    """
+    world_size = num_machines * num_gpus_per_machine
+    if world_size > 1:
+        if dist_url == "auto":
+            assert num_machines == 1, "dist_url=auto not supported in multi-machine jobs."
+            port = _find_free_port()
+            dist_url = f"tcp://127.0.0.1:{port}"
+        if num_machines > 1 and dist_url.startswith("file://"):
+            logger = logging.getLogger(__name__)
+            logger.warning("file:// is not a reliable init_method in multi-machine jobs. Prefer tcp://")
+
+        process_context = mp.spawn(
+            _distributed_worker,
+            nprocs=num_gpus_per_machine,
+            args=(
+                main_func,
+                world_size,
+                num_gpus_per_machine,
+                machine_rank,
+                dist_url,
+                config,
+                timeout,
+            ),
+            join=False,
+        )
+        try:
+            process_context.join()
+        except KeyboardInterrupt:
+            for i, process in enumerate(process_context.processes):
+                if process.is_alive():
+                    logger.info("Terminating process %s", str(i))
+                    process.terminate()
+                process.join()
+                logger.info("Process %s finished", str(i))
+    else:
+        main_func(config)
+
 
 @hydra.main(config_path="../configs", config_name="config.yaml")
 def main(config: DictConfig):
     """Main function."""
-    trainer = Trainer(config)
-    trainer.run_multiprocess_train()
+    launch(
+        _train,
+        config.machine_config.num_gpus,
+        num_machines=config.machine_config.num_machines,
+        machine_rank=config.machine_config.machine_rank,
+        dist_url=config.machine_config.dist_url,
+        config=(config,),
+    )
 
 
 if __name__ == "__main__":
diff --git a/tests/nerf/configs/test_basic_graph.yaml b/tests/nerf/configs/test_basic_graph.yaml
index 77939a274..96d5c56ba 100644
--- a/tests/nerf/configs/test_basic_graph.yaml
+++ b/tests/nerf/configs/test_basic_graph.yaml
@@ -1,4 +1,8 @@
-world_size: 0
+machine_config:
+  num_gpus: 8
+  num_machines: 1
+  machine_rank: 0
+  dist_url: auto
 network:
   mlp_0:
     class_name: MLP

From 62d4218421ef68bd7ec840c009f223757aa2f7dd Mon Sep 17 00:00:00 2001
From: evonneng <evonne_ng@berkeley.edu>
Date: Fri, 15 Apr 2022 21:53:35 +0000
Subject: [PATCH 3/6] updated logic within Trainer class to allow for
 multiprocess training

---
 mattport/nerf/trainer.py   | 19 ++++++++++++-------
 scripts/run_train_nerf.py  |  7 +++++--
 tests/nerf/test_trainer.py | 15 ---------------
 3 files changed, 17 insertions(+), 24 deletions(-)
 delete mode 100644 tests/nerf/test_trainer.py

diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index 549934573..0b900d6c8 100644
--- a/mattport/nerf/trainer.py
+++ b/mattport/nerf/trainer.py
@@ -1,28 +1,29 @@
 """
 Code to train model.
 """
-import os
-
-import numpy as np
-import torch
 import torch.distributed as dist
-import torch.multiprocessing as mp
-from omegaconf import DictConfig
+from torch.nn.parallel import DistributedDataParallel as DDP
 
 from mattport.nerf.dataset.build import build_dataset
+from mattport.nerf.graph import Graph
 from mattport.nerf.optimizer import Optimizer
 
 
 class Trainer:
     """Training class"""
 
-    def __init__(self, config: dict):
+    def __init__(self, local_rank: int, world_size: int, config: dict):
+        self.local_rank = local_rank
+        self.world_size = world_size
         self.config = config
         self.train_dataset = None
         self.test_dataset = None
         self.graph = None
         self.optimizer = None
 
+        self.device = f"cuda:{local_rank}"
+        self.is_main_thread = self.local_rank % self.world_size == 0
+
     def setup_dataset(self):
         """_summary_"""
         self.train_dataset = build_dataset(self.config.dataset)
@@ -30,6 +31,10 @@ class Trainer:
 
     def setup_graph(self):
         """_summary_"""
+        self.graph = Graph(self.config).to(self.device)
+        if self.world_size > 1:
+            self.graph = DDP(self.graph, device_ids=[self.local_rank])
+            dist.barrier(device_ids=[self.local_rank])
         raise NotImplementedError
 
     def setup_optimizer(self):
diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py
index 90a1ab246..54d37f099 100644
--- a/scripts/run_train_nerf.py
+++ b/scripts/run_train_nerf.py
@@ -112,8 +112,11 @@ def _train(local_rank: int, world_size: int, config: DictConfig) -> Any:
     Returns:
         Any: TODO(): determine the return type
     """
-    # trainer = Trainer(local_rank, world_size, config)
-    # trainer.train()
+    trainer = Trainer(local_rank, world_size, config)
+    trainer.setup_graph()
+    trainer.setup_optimizer()
+    trainer.setup_dataset()
+    trainer.train()
     return 0
 
 
diff --git a/tests/nerf/test_trainer.py b/tests/nerf/test_trainer.py
deleted file mode 100644
index 0cc9ff4dd..000000000
--- a/tests/nerf/test_trainer.py
+++ /dev/null
@@ -1,15 +0,0 @@
-from tty import CFLAG
-import pytest
-import runpy
-
-from mattport.nerf.trainer import Trainer
-from hydra import compose, initialize
-
-CONFIG_DIR = "configs"
-
-def test_multiprocess_train():
-    """Test for trainer to validate distributed data parallel working."""
-    with initialize(config_path=CONFIG_DIR):
-        cfg = compose(config_name='test_basic_graph')
-    trainer = Trainer(cfg)
-    trainer.run_multiprocess_train()
\ No newline at end of file

From b8e27136430bb9940a7041bad7dddba5ec1523c3 Mon Sep 17 00:00:00 2001
From: evonneng <evonne_ng@berkeley.edu>
Date: Fri, 15 Apr 2022 23:35:00 +0000
Subject: [PATCH 4/6] Changed Graph to nn.ModuleDict; Added DDP logic to
 seup_graph; locallyrunning DDP

---
 mattport/nerf/field_modules/mlp.py |  2 +-
 mattport/nerf/graph.py             | 11 +++++------
 mattport/nerf/trainer.py           |  6 +++---
 scripts/run_train_nerf.py          |  4 ++--
 tests/nerf/graph/test_graph.py     | 12 ++++++------
 5 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/mattport/nerf/field_modules/mlp.py b/mattport/nerf/field_modules/mlp.py
index 008920d97..57c99a900 100644
--- a/mattport/nerf/field_modules/mlp.py
+++ b/mattport/nerf/field_modules/mlp.py
@@ -39,7 +39,7 @@ class MLP(FieldModule):
         self.out_activation = out_activation
         self.net = None
 
-    def build_model(self) -> None:
+    def build_nn_modules(self) -> None:
         """Initialize mulilayer perceptron."""
         layers = []
         for i in range(self.num_layers):
diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py
index 5dd87440d..1f951e8bd 100644
--- a/mattport/nerf/graph.py
+++ b/mattport/nerf/graph.py
@@ -30,7 +30,7 @@ class Node:
         return hash(self.name)
 
 
-class Graph(nn.Module):
+class Graph(nn.ModuleDict):
     """_summary_"""
 
     def __init__(self, modules_config: DictConfig) -> None:
@@ -40,20 +40,19 @@ class Graph(nn.Module):
         self.roots = self.construct_graph()
         self.module_order = self.get_module_order()
 
-        self.modules = {}
         # initialize graph with known input dimensions; set default in_dim to 0
         for module_name, module_dict in modules_config.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.modules[module_name] = module(**module_dict.meta_data)
+            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.modules.items():
+        for _, module in self.items():
             module.build_nn_modules()
 
     def get_in_dim(self, curr_node: Node) -> None:
@@ -66,8 +65,8 @@ class Graph(nn.Module):
         if len(curr_node.parents) > 0:
             in_dim = 0
             for parent_name in curr_node.parents.keys():
-                in_dim += self.modules[parent_name].get_out_dim()
-            self.modules[curr_node.name].set_in_dim(in_dim)
+                in_dim += self[parent_name].get_out_dim()
+            self[curr_node.name].set_in_dim(in_dim)
             self.modules_config[curr_node.name].meta_data.in_dim = in_dim
 
         for child_node in curr_node.children.values():
diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index 0b900d6c8..de8ccafc3 100644
--- a/mattport/nerf/trainer.py
+++ b/mattport/nerf/trainer.py
@@ -1,6 +1,7 @@
 """
 Code to train model.
 """
+from omegaconf import DictConfig
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 
@@ -12,7 +13,7 @@ from mattport.nerf.optimizer import Optimizer
 class Trainer:
     """Training class"""
 
-    def __init__(self, local_rank: int, world_size: int, config: dict):
+    def __init__(self, local_rank: int, world_size: int, config: DictConfig):
         self.local_rank = local_rank
         self.world_size = world_size
         self.config = config
@@ -31,11 +32,10 @@ class Trainer:
 
     def setup_graph(self):
         """_summary_"""
-        self.graph = Graph(self.config).to(self.device)
+        self.graph = Graph(self.config.network).to(self.device)
         if self.world_size > 1:
             self.graph = DDP(self.graph, device_ids=[self.local_rank])
             dist.barrier(device_ids=[self.local_rank])
-        raise NotImplementedError
 
     def setup_optimizer(self):
         """_summary_"""
diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py
index 54d37f099..e0e4dbbf2 100644
--- a/scripts/run_train_nerf.py
+++ b/scripts/run_train_nerf.py
@@ -174,7 +174,7 @@ def launch(
                 process.join()
                 logger.info("Process %s finished", str(i))
     else:
-        main_func(config)
+        main_func(local_rank=0, world_size=1, config=config)
 
 
 @hydra.main(config_path="../configs", config_name="config.yaml")
@@ -186,7 +186,7 @@ def main(config: DictConfig):
         num_machines=config.machine_config.num_machines,
         machine_rank=config.machine_config.machine_rank,
         dist_url=config.machine_config.dist_url,
-        config=(config,),
+        config=config,
     )
 
 
diff --git a/tests/nerf/graph/test_graph.py b/tests/nerf/graph/test_graph.py
index e45665341..191ae9aa8 100644
--- a/tests/nerf/graph/test_graph.py
+++ b/tests/nerf/graph/test_graph.py
@@ -32,17 +32,17 @@ def instantiate_graph(config_name: str) -> Graph:
 def test_graph_init():
     """test field modules dictionary is instantiated properly"""
     test_graph, _ = instantiate_graph("test_basic_graph")
-    assert test_graph.modules is not None
-    assert type(test_graph.modules["encoder_0"]).__name__ == "ScalingAndOffset"
-    assert type(test_graph.modules["mlp_0"]).__name__ == "MLP"
-    assert type(test_graph.modules["mlp_1"]).__name__ == "MLP"
+    assert test_graph is not None
+    assert type(test_graph["encoder_0"]).__name__ == "ScalingAndOffset"
+    assert type(test_graph["mlp_0"]).__name__ == "MLP"
+    assert type(test_graph["mlp_1"]).__name__ == "MLP"
 
 
 def test_input_dimension():
     """test field module input dimension calculation"""
     test_graph, _ = instantiate_graph("test_repeat_graph")
-    assert test_graph.modules["mlp_0"].in_dim == 16
-    assert test_graph.modules["mlp_1"].in_dim == 1
+    assert test_graph["mlp_0"].in_dim == 16
+    assert test_graph["mlp_1"].in_dim == 1
 
 
 def test_construct_graph_basic():

From 563dba0bad789b4a20660c94330c1a448f075464 Mon Sep 17 00:00:00 2001
From: evonneng <evonne_ng@berkeley.edu>
Date: Sat, 16 Apr 2022 01:25:55 +0000
Subject: [PATCH 5/6] lint

---
 mattport/nerf/trainer.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index de8ccafc3..3c80c884b 100644
--- a/mattport/nerf/trainer.py
+++ b/mattport/nerf/trainer.py
@@ -22,9 +22,6 @@ class Trainer:
         self.graph = None
         self.optimizer = None
 
-        self.device = f"cuda:{local_rank}"
-        self.is_main_thread = self.local_rank % self.world_size == 0
-
     def setup_dataset(self):
         """_summary_"""
         self.train_dataset = build_dataset(self.config.dataset)
@@ -32,7 +29,7 @@ class Trainer:
 
     def setup_graph(self):
         """_summary_"""
-        self.graph = Graph(self.config.network).to(self.device)
+        self.graph = Graph(self.config.network).to(f"cuda:{self.local_rank}")
         if self.world_size > 1:
             self.graph = DDP(self.graph, device_ids=[self.local_rank])
             dist.barrier(device_ids=[self.local_rank])
@@ -59,6 +56,7 @@ class Trainer:
 
     def train(self) -> None:
         """_summary_"""
+        # is_main_thread = self.local_rank % self.world_size == 0
         raise NotImplementedError
 
     def train_epoch(self):

From c85ee76b1f8bb9d7eee2c99f8a4862ee04b74c0b Mon Sep 17 00:00:00 2001
From: Ethan Weber <eweb0124@gmail.com>
Date: Sun, 17 Apr 2022 16:45:38 +0000
Subject: [PATCH 6/6] removed build_dataset

---
 mattport/nerf/trainer.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index b65f01062..6d339404b 100644
--- a/mattport/nerf/trainer.py
+++ b/mattport/nerf/trainer.py
@@ -5,7 +5,6 @@ from omegaconf import DictConfig
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 
-from mattport.nerf.dataset.build import build_dataset
 from mattport.nerf.graph import Graph
 from mattport.nerf.optimizer import Optimizer