From bf3b77f952bcc6bf9e48dd90829209561db598c5 Mon Sep 17 00:00:00 2001 From: evonneng Date: Wed, 20 Apr 2022 12:06:29 -0700 Subject: [PATCH] Eng/distributed training (#27) * added checkpointing logic * small nits --- mattport/nerf/trainer.py | 62 +++++++++++++++++++++++++----------- mattport/utils/decorators.py | 14 ++++++++ mattport/utils/writer.py | 17 +++------- 3 files changed, 62 insertions(+), 31 deletions(-) create mode 100644 mattport/utils/decorators.py diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index ebd27ccc3..bb4e39d64 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -1,8 +1,10 @@ """ Code to train model. """ - import logging +import os + +import torch import torch.distributed as dist from hydra.utils import instantiate from omegaconf import DictConfig @@ -13,6 +15,7 @@ from tqdm import tqdm 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 class Trainer: @@ -30,6 +33,7 @@ class Trainer: self.test_dataloader = None self.graph = None self.optimizer = None + self.is_main_thread = self.local_rank % self.world_size == 0 def setup_dataset(self): """_summary_""" @@ -54,6 +58,9 @@ class Trainer: self.graph = instantiate( self.config.graph, intrinsics=dataset_inputs.intrinsics, camera_to_world=dataset_inputs.camera_to_world ).to(f"cuda:{self.local_rank}") + + ## TODO(): we need to call load_checkpoint here if need be (before we wrap it in DDP) + if self.world_size > 1: self.graph = DDP(self.graph, device_ids=[self.local_rank]) dist.barrier(device_ids=[self.local_rank]) @@ -62,26 +69,46 @@ class Trainer: """_summary_""" self.optimizer = optim.Adam(self.graph.parameters(), lr=self.config.optimizer.lr) - def load_checkpoint(self): - """_summary_ + def load_checkpoint(self, load_path: str) -> int: + """Load the checkpoint from the given path - Raises: - NotImplementedError: _description_ + Args: + load_path (str): path from which to load the model + + Returns: + int: step iteration associated with the loaded checkpoint """ - raise NotImplementedError + assert os.path.exists(load_path), f"Checkpoint {load_path} does not exist" + loaded_state = torch.load(load_path, map_location="cpu") + self.graph.load_state_dict({key.replace("module.", ""): value for key, value in loaded_state["model"].items()}) + self.optimizer.optimizer.load_state_dict(loaded_state["optimizer"]) + step = loaded_state["step"] + logging.info("done loading checkpoint from %s", load_path) + return step - def save_checkpoint(self): - """_summary_ + @check_main_thread + def save_checkpoint(self, output_dir: str, step: int) -> None: + """Save the model and optimizer - Raises: - NotImplementedError: _description_ + Args: + output_dir (str): directory to save the checkpoint + step (int): number of steps in training for given checkpoint + model (Graph): Graph model to be saved + optimizer (Optimizer): Optimizer to be saved """ - raise NotImplementedError + ckpt_path = os.path.join(output_dir, f"step-{step:09d}.ckpt") + torch.save( + { + "step": step, + "model": self.graph.module.state_dict() if hasattr(self.graph, "module") else self.graph.state_dict(), + "optimizer": self.optimizer.optimizer.state_dict(), + }, + ckpt_path, + ) + logging.info("done saving checkpoint to %s", ckpt_path) def train(self) -> None: """_summary_""" - # is_main_thread = self.local_rank % self.world_size == 0 - # raise NotImplementedError num_epochs = 10 for _ in range(num_epochs): self.train_epoch() @@ -89,16 +116,15 @@ class Trainer: def train_epoch(self): """_summary_""" num_iters = 100 - for _ in tqdm(range(num_iters)): + for step in tqdm(range(num_iters)): batch = next(iter(self.train_dataloader)) self.train_iteration(batch) + # TODO(): proper saving with the correct directory path and such + # if step % self.config.save_step == 0: + # self.save_checkpoint(self.config.model_dir) def train_iteration(self, batch): """_summary_""" - # TODO(): save checkpoints, and do logging only on rank 0 device only - if self.local_rank == 0: - pass - # move batch to correct device ray_indices = batch.indices.to(f"cuda:{self.local_rank}") graph_outputs = self.graph(ray_indices) diff --git a/mattport/utils/decorators.py b/mattport/utils/decorators.py new file mode 100644 index 000000000..7a016996a --- /dev/null +++ b/mattport/utils/decorators.py @@ -0,0 +1,14 @@ +""" +Decorator definitions +""" +from typing import Callable + + +def check_main_thread(func: Callable) -> Callable: + """Declarator: check if you are on main thread""" + + def wrapper(self, *args, **kwargs): + if self.is_main_thread: + func(self, *args, **kwargs) + + return wrapper diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index 76405f4c5..ba8011cf0 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -4,21 +4,12 @@ Generic Writer class from abc import abstractmethod -from typing import Callable import torch from torch.utils.tensorboard import SummaryWriter from torchtyping import TensorType - -def require_main_thread(func: Callable) -> Callable: - """Declarator: check if you are on main thread""" - - def wrapper(self, *args, **kwargs): - if self.is_main_thread: - func(self, *args, **kwargs) - - return wrapper +from mattport.utils.decorators import check_main_thread class Writer: @@ -79,7 +70,7 @@ class TensorboardWriter(Writer): if self.is_main_thread: self.writer = SummaryWriter(log_dir=self.save_dir) - @require_main_thread + @check_main_thread def write_image(self, name: str, x: TensorType["H", "W", 3]) -> None: """_summary_ @@ -91,7 +82,7 @@ class TensorboardWriter(Writer): x = to8b(x) self.writer.add_images(name, x) - @require_main_thread + @check_main_thread def write_text(self, name: str, x: str) -> None: """Tensorboard method to write a string to summary @@ -101,7 +92,7 @@ class TensorboardWriter(Writer): """ self.writer.add_text(name, x) - @require_main_thread + @check_main_thread def write_scalar(self, name: str, x: float, y: float) -> None: """Tensorboard method to write a single scalar value to the logger