From 4d1a6b2dbd69e9a676a5a0f4d06f9e1be2395142 Mon Sep 17 00:00:00 2001 From: evonneng Date: Thu, 21 Apr 2022 11:39:44 -0700 Subject: [PATCH] cleanup and added logic to train for load/save checkpoints (#32) * added checkpointing logic * small nits * check DDP running correctly * cleanup and added logic to train for load/save checkpoints * lint --- configs/config.yaml | 8 +++++ mattport/nerf/graph/base.py | 66 +------------------------------------ mattport/nerf/trainer.py | 57 +++++++++++++++++--------------- mattport/utils/writer.py | 4 +-- scripts/run_train_nerf.py | 1 - 5 files changed, 42 insertions(+), 94 deletions(-) diff --git a/configs/config.yaml b/configs/config.yaml index 85f57247b..9e2847eea 100644 --- a/configs/config.yaml +++ b/configs/config.yaml @@ -1,3 +1,11 @@ +# pretrained loading options +resume_train: + load_dir: /tmp/mattport_models/ + load_step: 10 + +save_step: 10 +model_dir: /tmp/mattport_models/ + # dataset options dataset: data_directory: data/lego diff --git a/mattport/nerf/graph/base.py b/mattport/nerf/graph/base.py index 750dd3072..7da490633 100644 --- a/mattport/nerf/graph/base.py +++ b/mattport/nerf/graph/base.py @@ -1,10 +1,9 @@ """ The Graph module contains all trainable parameters. """ -import importlib from abc import abstractmethod from dataclasses import dataclass -from typing import Dict, List, Optional, Set +from typing import Dict, Optional from torch import nn from torchtyping import TensorType @@ -62,69 +61,6 @@ class Graph(nn.Module): if not child_node.visited_in_dim: self.get_in_dim(child_node) - def construct_graph(self) -> Set[Node]: - """Constructs a dependency graph given the module configuration - - Args: - config (dict): module definitions that make up the network - - Returns: - set: all root nodes of the constructed dependency graph - """ - processed_modules = {} - roots = set() - for module_name, module_dict in self.modules.items(): - if not module_name in processed_modules: - curr_module = Node(name=module_name, children={}, parents={}) - processed_modules[module_name] = curr_module - else: - curr_module = processed_modules[module_name] - inputs = module_dict.inputs - for input_module in inputs: - if not input_module in processed_modules: - parent_module = Node( - name=input_module, - children={module_name: curr_module}, - parents={}, - ) - processed_modules[input_module] = parent_module - else: - processed_modules[input_module].children[module_name] = curr_module - if input_module == "x": - roots.add(curr_module) - else: - curr_module.parents[input_module] = parent_module - return roots - - def topological_sort(self, curr_node: Node, ordering_stack: List[str]) -> None: - """utility function to sort the call order in the dependency graph - - Args: - curr_node (Node): pointer to current node in process - ordering_stack (list): cumulative ordering of graph nodes - """ - curr_node.visited_order = True - for child_node in curr_node.children.values(): - if not child_node.visited_order: - self.topological_sort(child_node, ordering_stack) - ordering_stack.append(curr_node.name) - - def get_module_order(self) -> List[str]: - """Generates a graph and determines order of module operations using topological sorting - - Args: - config (dict): module definitions that make up the network - - Returns: - list: ordering of the module names that should be executed - """ - roots = self.roots - ordering_stack = [] - for root in roots: - if not root.visited_order: - self.topological_sort(root, ordering_stack) - return ordering_stack[::-1] - @abstractmethod def forward(self, ray_indices: TensorType["num_rays", 3]): """Forward function that needs to be overridden.""" diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 74e7d01f5..a8e866e98 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -11,13 +11,14 @@ from omegaconf import DictConfig from torch import optim from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader -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 +from mattport.utils.writer import LocalWriter + +logging.getLogger("PIL").setLevel(logging.WARNING) -logging.getLogger('PIL').setLevel(logging.WARNING) class Trainer: """Training class""" @@ -28,13 +29,18 @@ class Trainer: self.config = config self.local_rank = local_rank self.world_size = world_size + # dataset variables self.train_dataset = None self.train_dataloader = None self.test_dataset = None self.test_dataloader = None + # model variables self.graph = None self.optimizer = None - self.is_main_thread = self.local_rank % self.world_size == 0 + self.start_step = 0 + # logging variables + self.is_main_thread = local_rank % world_size == 0 + self.local_writer = LocalWriter(local_rank, world_size, save_dir=None) def setup_dataset(self): """_summary_""" @@ -59,8 +65,10 @@ 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}") + self.setup_optimizer() - ## TODO(): we need to call load_checkpoint here if need be (before we wrap it in DDP) + if self.config.resume_train: + self.load_checkpoint(self.config.resume_train) if self.world_size > 1: self.graph = DDP(self.graph, device_ids=[self.local_rank]) @@ -70,7 +78,7 @@ class Trainer: """_summary_""" self.optimizer = optim.Adam(self.graph.parameters(), lr=self.config.optimizer.lr) - def load_checkpoint(self, load_path: str) -> int: + def load_checkpoint(self, load_config: DictConfig) -> int: """Load the checkpoint from the given path Args: @@ -79,13 +87,13 @@ class Trainer: Returns: int: step iteration associated with the loaded checkpoint """ + load_path = os.path.join(load_config.load_dir, f"step-{load_config.load_step:09d}.ckpt") 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"] + self.optimizer.load_state_dict(loaded_state["optimizer"]) + self.start_step = loaded_state["step"] + 1 logging.info("done loading checkpoint from %s", load_path) - return step @check_main_thread def save_checkpoint(self, output_dir: str, step: int) -> None: @@ -97,38 +105,35 @@ class Trainer: model (Graph): Graph model to be saved optimizer (Optimizer): Optimizer to be saved """ + if not os.path.exists(output_dir): + os.makedirs(output_dir) 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(), + "optimizer": self.optimizer.state_dict(), }, ckpt_path, ) - logging.info("done saving checkpoint to %s", ckpt_path) def train(self) -> None: """_summary_""" num_epochs = 10 for _ in range(num_epochs): - self.train_epoch() + self.train_iteration() - def train_epoch(self): + def train_iteration(self): """_summary_""" num_iters = 100 - for _ in tqdm(range(num_iters)): + for step in range(self.start_step, self.start_step + 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_""" - # move batch to correct device - 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.module.get_losses(batch, graph_outputs) - # logging.info(losses) + # move batch to correct device + 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.module.get_losses(batch, graph_outputs) + self.local_writer.write_text(f"{step}/{num_iters}: losses", losses) + if step % self.config.save_step == 0: + self.save_checkpoint(self.config.model_dir, step) + self.local_writer.write_text(f"ckpt at {step}", f"{losses}\n") diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index 7fb65d89f..3b7dc7c79 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -3,12 +3,12 @@ Generic Writer class """ -from abc import abstractmethod import os +from abc import abstractmethod + import imageio import numpy as np import torch - from torch.utils.tensorboard import SummaryWriter from torchtyping import TensorType diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py index a71150041..b7f0835c4 100644 --- a/scripts/run_train_nerf.py +++ b/scripts/run_train_nerf.py @@ -114,7 +114,6 @@ def _train(local_rank: int, world_size: int, config: DictConfig) -> Any: """ trainer = Trainer(config, local_rank, world_size) trainer.setup_graph() - trainer.setup_optimizer() trainer.setup_dataset() trainer.train() return 0