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 66ec3b237..6d339404b 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -1,13 +1,20 @@ """ Code to train model. """ +from omegaconf import DictConfig +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP + +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: DictConfig): + self.local_rank = local_rank + self.world_size = world_size self.config = config self.train_dataset = None self.test_dataset = None @@ -21,7 +28,10 @@ class Trainer: def setup_graph(self): """_summary_""" - raise NotImplementedError + 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]) def setup_optimizer(self): """_summary_""" @@ -43,8 +53,9 @@ class Trainer: """ raise NotImplementedError - def train(self): + def train(self) -> None: """_summary_""" + # is_main_thread = self.local_rank % self.world_size == 0 raise NotImplementedError def train_epoch(self): @@ -53,4 +64,6 @@ class Trainer: def train_iteration(self): """_summary_""" + # TODO(): save checkpoints, and do logging only on rank 0 device only + # if gpu == 0: raise NotImplementedError diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py index 2ad9a5f97..e0e4dbbf2 100644 --- a/scripts/run_train_nerf.py +++ b/scripts/run_train_nerf.py @@ -1,23 +1,193 @@ """ 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.setup_graph() + trainer.setup_optimizer() + trainer.setup_dataset() + 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(local_rank=0, world_size=1, config=config) + @hydra.main(config_path="../configs", config_name="config.yaml") def main(config: DictConfig): """Main function.""" - trainer = Trainer(config) - trainer.setup_dataset() - trainer.setup_graph() - trainer.setup_optimizer() + 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/graph/configs/test_basic_graph.yaml b/tests/nerf/graph/configs/test_basic_graph.yaml index b3e23bd79..96d5c56ba 100644 --- a/tests/nerf/graph/configs/test_basic_graph.yaml +++ b/tests/nerf/graph/configs/test_basic_graph.yaml @@ -1,3 +1,8 @@ +machine_config: + num_gpus: 8 + num_machines: 1 + machine_rank: 0 + dist_url: auto network: mlp_0: class_name: MLP 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():