From a498a6cc64831efa49caed22fab0cc274b70d7b0 Mon Sep 17 00:00:00 2001 From: evonneng Date: Fri, 6 May 2022 21:55:33 -0700 Subject: [PATCH] Standardized logging code to all global var; writers now accessible from all files (#62) * added skeleton for terminal logger * done skeleton terminal writer logic... todo: actual calls to the write classes * moved lambda to top function * added logic to support timing * prettier printing * prettier printing * lint * separated out timing and other stats variables. printing in dict format * keeping logging history and wiping rest * updated logic to calculate total times instead of average times * made stats tracker configurable * moved stats tracker and profiler to own class file; implemented basic functionality for profiler * lint removing duplicate function * update readme; moved around config variables * profiler prints on sig kill * merge * standardized all logging to global variables --- configs/logging/default_logging.yaml | 4 +- mattport/nerf/trainer.py | 70 ++++++++++++++-------- mattport/utils/comms.py | 72 +++++++++++++++++++++++ mattport/utils/decorators.py | 6 +- mattport/utils/profiler.py | 27 +++++++-- mattport/utils/stats_tracker.py | 46 ++++++++++++--- mattport/utils/writer.py | 50 ++++++++++++---- scripts/run_train.py | 26 ++++---- tests/configs/test_default.yml | 88 ++++++++++++++++++++++++++++ 9 files changed, 322 insertions(+), 67 deletions(-) create mode 100644 mattport/utils/comms.py create mode 100644 tests/configs/test_default.yml diff --git a/configs/logging/default_logging.yaml b/configs/logging/default_logging.yaml index 1e74b13c4..59ef1333b 100644 --- a/configs/logging/default_logging.yaml +++ b/configs/logging/default_logging.yaml @@ -6,7 +6,7 @@ stats_tracker: stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME, RAYS_PER_SEC, CURR_TEST_PSNR] # see mattport/utils/stats_tracker.py for options # writer logs losses and images per iteration writer: - type: TensorboardWriter - save_dir: './' + TensorboardWriter: + save_dir: './' # profiler logs run times of functions and prints at end of training enable_profiler: True \ No newline at end of file diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index de289a29f..e4e9a938c 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -13,15 +13,13 @@ from omegaconf import DictConfig from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader -import mattport.utils.writer from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch from mattport.nerf.dataset.utils import DatasetInputs, get_dataset_inputs_dict from mattport.nerf.metrics import get_psnr from mattport.nerf.optimizers import Optimizers -from mattport.utils import profiler +from mattport.utils import profiler, stats_tracker, writer from mattport.utils.decorators import check_main_thread -from mattport.utils.stats_tracker import Stats, StatsTracker logging.getLogger("PIL").setLevel(logging.WARNING) @@ -51,12 +49,9 @@ class Trainer: self.optimizers = None self.start_step = 0 # logging variables - self.is_main_thread = world_size != 0 and local_rank % world_size == 0 - writer = getattr(mattport.utils.writer, self.config.logging.writer.type) - self.writer = writer(self.is_main_thread, save_dir=self.config.logging.writer.save_dir) - self.stats = StatsTracker(config, self.is_main_thread) - if not profiler.PROFILER and self.config.logging.enable_profiler: - profiler.PROFILER = profiler.Profiler(config, self.is_main_thread) + stats_tracker.setup_stats_tracker(config) + writer.setup_event_writers(config.logging.writer) + profiler.setup_profiler(config.logging) self.device = "cpu" if self.world_size == 0 else f"cuda:{self.local_rank}" @profiler.time_function @@ -187,27 +182,44 @@ class Trainer: for i, step in enumerate(range(self.start_step, self.start_step + num_iterations)): data_start = time() batch = next(iter_dataset) - self.stats.update_time(Stats.ITER_LOAD_TIME, data_start, time(), step=step) + stats_tracker.update_stats( + {"name": stats_tracker.Stats.ITER_LOAD_TIME, "start_time": data_start, "end_time": time(), "step": step} + ) iter_start = time() loss_dict = self.train_iteration(batch, step) - self.stats.update_time( - Stats.RAYS_PER_SEC, iter_start, time(), step=step, batch_size=batch["indices"].shape[0] + stats_tracker.update_stats( + { + "name": stats_tracker.Stats.RAYS_PER_SEC, + "start_time": iter_start, + "end_time": time(), + "step": step, + "batch_size": batch["indices"].shape[0], + }, + ) + stats_tracker.update_stats( + { + "name": stats_tracker.Stats.ITER_TRAIN_TIME, + "start_time": iter_start, + "end_time": time(), + "step": step, + } ) - self.stats.update_time(Stats.ITER_TRAIN_TIME, iter_start, time(), step=step) if step != 0 and step % self.config.logging.steps_per_log == 0: - self.writer.write_scalar_dict(loss_dict, step, group="Loss", prefix="train-") + writer.write_event({"scalar_dict": loss_dict, "step": step, "group": "Loss", "prefix": "train-"}) # TODO: add the learning rates to tensorboard/logging if step != 0 and self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0: self.save_checkpoint(self.config.graph.model_dir, step) if step % self.config.graph.steps_per_test == 0: for image_idx in self.config.data.validation_image_indices: self.test_image(image_idx=image_idx, step=step) - self.stats.print_stats(i / num_iterations) + stats_tracker.print_stats(i / num_iterations) - self.stats.update_time(Stats.TOTAL_TRAIN_TIME, train_start, time(), step=-1) # NOTE(ethan): why is step -1? - self.stats.print_stats(-1) + stats_tracker.update_stats( + {"name": stats_tracker.Stats.TOTAL_TRAIN_TIME, "start_time": train_start, "end_time": time()} + ) + stats_tracker.print_stats(1.0) @profiler.time_function def train_iteration(self, batch: dict, step: int): @@ -266,17 +278,29 @@ class Trainer: disparity_fine = torch.cat(disparity_fine).view(image_height, image_width, 1).detach().cpu() combined_image = torch.cat([image, rgb_coarse, rgb_fine], dim=1) - self.writer.write_image(f"image_idx_{image_idx}-rgb_coarse_fine", combined_image, step, group="val_img") + writer.write_event( + {"name": f"image_idx_{image_idx}-rgb_coarse_fine", "x": combined_image, "step": step, "group": "val_img"} + ) combined_image = torch.cat([accumulation_coarse, accumulation_fine], dim=1) - self.writer.write_image(f"image_idx_{image_idx}", combined_image, step, group="val_accumulation") + writer.write_event( + {"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_accumulation"} + ) combined_image = torch.cat([disparity_coarse, disparity_fine], dim=1) - self.writer.write_image(f"image_idx_{image_idx}", combined_image, step, group="val_disparity") + writer.write_event( + {"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_disparity"} + ) coarse_psnr = get_psnr(image, rgb_coarse) - self.writer.write_scalar(f"image_idx_{image_idx}", float(coarse_psnr), step, group="val") + writer.write_event( + {"name": f"image_idx_{image_idx}", "scalar": float(coarse_psnr), "step": step, "group": "val"} + ) fine_psnr = get_psnr(image, rgb_fine) - self.stats.update_value(Stats.CURR_TEST_PSNR, float(fine_psnr), step) - self.writer.write_scalar(f"image_idx_{image_idx}-fine_psnr", float(fine_psnr), step, group="val") + stats_tracker.update_stats( + {"name": stats_tracker.Stats.CURR_TEST_PSNR, "value": float(fine_psnr), "step": step} + ) + writer.write_event( + {"name": f"image_idx_{image_idx}-fine_psnr", "scalar": float(fine_psnr), "step": step, "group": "val"} + ) diff --git a/mattport/utils/comms.py b/mattport/utils/comms.py new file mode 100644 index 000000000..9a844ec2e --- /dev/null +++ b/mattport/utils/comms.py @@ -0,0 +1,72 @@ +"""functionality to handle multiprocessing syncing and communicating""" +import torch +import torch.distributed as dist + +_LOCAL_PROCESS_GROUP = None + + +def get_world_size() -> int: + """Get total number of available gpus""" + if not dist.is_available(): + return 1 + if not dist.is_initialized(): + return 1 + return dist.get_world_size() + + +def get_rank() -> int: + """Get global rank of current thread""" + if not dist.is_available(): + return 0 + if not dist.is_initialized(): + return 0 + return dist.get_rank() + + +def get_local_rank() -> int: + """The rank of the current process within the local (per-machine) process group.""" + if not dist.is_available(): + return 0 + if not dist.is_initialized(): + return 0 + assert ( + _LOCAL_PROCESS_GROUP is not None + ), "Local process group is not created! Please use launch() to spawn processes!" + return dist.get_rank(group=_LOCAL_PROCESS_GROUP) + + +def get_local_size() -> int: + """ + The size of the per-machine process group, + i.e. the number of processes per machine. + """ + if not dist.is_available(): + return 1 + if not dist.is_initialized(): + return 1 + return dist.get_world_size(group=_LOCAL_PROCESS_GROUP) + + +def is_main_process() -> bool: + """check to see if you are currently on the main process""" + return get_rank() == 0 + + +def synchronize(world_size): + """ + Helper function to synchronize (barrier) among all processes when + using distributed training + """ + if not dist.is_available(): + return + if not dist.is_initialized(): + return + world_size = dist.get_world_size() + if world_size == 1: + return + if dist.get_backend() == dist.Backend.NCCL: + # This argument is needed to avoid warnings. + # It's valid only for NCCL backend. + dist.barrier(device_ids=[torch.cuda.current_device()]) + else: + dist.barrier() diff --git a/mattport/utils/decorators.py b/mattport/utils/decorators.py index 988ae7c56..f7b79f7bf 100644 --- a/mattport/utils/decorators.py +++ b/mattport/utils/decorators.py @@ -3,6 +3,8 @@ Decorator definitions """ from typing import Callable, List +from mattport.utils import comms + def decorate_all(decorators: List[Callable]) -> Callable: """A decorator to decorate all member functions of a class""" @@ -34,7 +36,7 @@ def check_profiler_enabled(func: Callable) -> Callable: def wrapper(self, *args, **kwargs): ret = None - if self.config.logging.enable_stats: + if self.config.enable_stats: ret = func(self, *args, **kwargs) return ret @@ -46,7 +48,7 @@ def check_main_thread(func: Callable) -> Callable: def wrapper(self, *args, **kwargs): ret = None - if self.is_main_thread: + if comms.is_main_process(): ret = func(self, *args, **kwargs) return ret diff --git a/mattport/utils/profiler.py b/mattport/utils/profiler.py index 0d7f12619..2b75296b9 100644 --- a/mattport/utils/profiler.py +++ b/mattport/utils/profiler.py @@ -4,13 +4,14 @@ Profiler base class and functionality import logging import sys import time -from typing import Callable +from typing import Callable, Dict from omegaconf import DictConfig +from mattport.utils import comms from mattport.utils.decorators import check_main_thread, check_profiler_enabled, decorate_all -PROFILER = None +PROFILER = [] def time_function(func: Callable) -> Callable: @@ -25,20 +26,34 @@ def time_function(func: Callable) -> Callable: for attr in func.__qualname__.split(".")[:-1]: class_str += f"{vals[attr].__qualname__}_" class_str += func.__name__ - PROFILER.update_time(class_str, start, time.time()) + PROFILER[0].update_time(class_str, start, time.time()) return ret return wrapper +def flush_profiler(config: Dict): + """Method that checks if profiler is enabled before flushing + + Args: + config (Dict): config check + """ + if config.enable_profiler and PROFILER: + PROFILER[0].print_profile() + + +def setup_profiler(config: DictConfig): + """Initialization of profilers""" + if comms.is_main_process(): + PROFILER.append(Profiler(config)) + + @decorate_all([check_profiler_enabled, check_main_thread]) class Profiler: """Profiler class""" - def __init__(self, config: DictConfig, is_main_thread: bool): + def __init__(self, config: DictConfig): self.config = config - if self.config.logging.enable_profiler: - self.is_main_thread = is_main_thread self.profiler_dict = {} def update_time(self, func_name: str, start_time: float, end_time: float): diff --git a/mattport/utils/stats_tracker.py b/mattport/utils/stats_tracker.py index 58a181658..aa5dfabaa 100644 --- a/mattport/utils/stats_tracker.py +++ b/mattport/utils/stats_tracker.py @@ -4,10 +4,38 @@ Stats tracker base class and functionality import datetime import enum import logging +from typing import Any, Dict from omegaconf import DictConfig + +from mattport.utils import comms from mattport.utils.decorators import check_main_thread, check_print_stats_step, check_stats_enabled, decorate_all +STATS_TRACKER = [] + + +def update_stats(args: Dict[str, Any]): + """update global stats tracker based on arguments""" + assert comms.is_main_process(), "Writing out with process other than main" + if "value" in args: + STATS_TRACKER[0].update_value(**args) + elif "start_time" in args: + STATS_TRACKER[0].update_time(**args) + else: + raise NotImplementedError + + +def print_stats(fraction_done: float): + """print current statistics information""" + assert comms.is_main_process(), "Writing out with process other than main" + STATS_TRACKER[0].print_stats(fraction_done) + + +def setup_stats_tracker(config: DictConfig): + """Initialization of stats tracker""" + if comms.is_main_process(): + STATS_TRACKER.append(StatsTracker(config)) + class Stats(enum.Enum): """Possible Stats values for StatsTracker @@ -28,9 +56,8 @@ class Stats(enum.Enum): class StatsTracker: """Stats Tracker class""" - def __init__(self, config: DictConfig, is_main_thread: bool): + def __init__(self, config: DictConfig): self.config = config - self.is_main_thread = is_main_thread if self.config.logging.enable_stats: self.max_history = self.config.logging.stats_tracker.max_history self.step = 0 @@ -57,14 +84,17 @@ class StatsTracker: self.new_key = not name in self.stats_dict or self.new_key self.stats_dict[name] = value - def update_time(self, name: enum.Enum, start_time: float, end_time: float, step: int, batch_size: int = None): + def update_time( + self, name: enum.Enum, start_time: float, end_time: float, step: int = None, batch_size: int = None + ): """update the stats dictionary with running averages/cumulative durations Args: name (enum.Enum): Enum name of statistic we are logging start_time (float): start time for the call in seconds end_time (float): end time when the call finished executing in seconds - step (int): number of total iteration steps. + step (int): number of total iteration steps. Defaults to None. + if None, reports duration without averaging batch_size (int, optional): total number of rays in a batch; if None, reports duration instead of batch per second. Defaults to None. """ @@ -76,12 +106,12 @@ class StatsTracker: # calculate the batch per second stat val = batch_size / val - if step == -1: - # logging total time instead of average - self.stats_dict[name] = val - else: + if step: # calculate updated average self.stats_dict[name] = (self.stats_dict.get(name, 0) * step + val) / (step + 1) + else: + # logging total time instead of average + self.stats_dict[name] = val if name == Stats.ITER_TRAIN_TIME and Stats.ETA in self.stats_to_track: # update ETA if logging iteration train time diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index b73068b6a..bee18aeba 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -5,17 +5,44 @@ Generic Writer class import os from abc import abstractmethod -from typing import Dict +from typing import Any, Dict import imageio import numpy as np import torch +from omegaconf import DictConfig from torch.utils.tensorboard import SummaryWriter from torchtyping import TensorType +import mattport.utils.writer +from mattport.utils import comms from mattport.utils.decorators import check_main_thread, decorate_all to8b = lambda x: (255 * torch.clamp(x, min=0, max=1)).to(torch.uint8) +EVENT_WRITERS = [] + + +def write_event(args: Dict[str, Any]): + """write out values for all the event writers""" + assert comms.is_main_process(), "Writing out with process other than main" + for writer in EVENT_WRITERS: + if "x" in args: + print(args["x"].shape) + writer.write_image(**args) + elif "scalar" in args: + writer.write_scalar(**args) + elif "scalar_dict" in args: + writer.write_scalar_dict(**args) + else: + raise NotImplementedError + + +def setup_event_writers(config: DictConfig) -> None: + """Initialization of all event writers specified in config""" + if comms.is_main_process(): + for writer_type in config: + curr_writer = getattr(mattport.utils.writer, writer_type) + EVENT_WRITERS.append(curr_writer(config[writer_type].save_dir)) def get_tensorboard_name(name: str, group: str = None, prefix: str = None): @@ -31,19 +58,18 @@ def get_tensorboard_name(name: str, group: str = None, prefix: str = None): class Writer: """Writer class""" - def __init__(self, is_main_thread: bool, save_dir: str): - self.is_main_thread = is_main_thread + def __init__(self, save_dir: str): self.save_dir = save_dir @abstractmethod def write_image( - self, name: str, x: TensorType["H", "W", 3], step: int, group: str = None, prefix: str = None + self, name: str, x: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None ) -> None: """_summary_ Args: name (str): data identifier - x (TensorType["H", "W", 3]): rendered image to write + x (TensorType["H", "W", "C"]): rendered image to write step (int): the time step to log group (str): the group e.g., "Loss", "Accuracy", "Time" prefix (str): the prefix e.g., "train-", "test-" @@ -74,19 +100,17 @@ class Writer: group (str): the group e.g., "Loss", "Accuracy", "Time" prefix (str): the prefix e.g., "train-", "test-" """ - if self.is_main_thread: - for name, scalar in scalar_dict.items(): - self.write_scalar(name, scalar, step, group=group, prefix=prefix) + for name, scalar in scalar_dict.items(): + self.write_scalar(name, scalar, step, group=group, prefix=prefix) @decorate_all([check_main_thread]) class TensorboardWriter(Writer): """Tensorboard Writer Class""" - def __init__(self, is_main_thread: bool, save_dir: str): - super().__init__(is_main_thread, save_dir) - if self.is_main_thread: - self.tb_writer = SummaryWriter(log_dir=self.save_dir) + def __init__(self, save_dir: str): + super().__init__(save_dir) + self.tb_writer = SummaryWriter(log_dir=self.save_dir) def write_image( self, name: str, x: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None @@ -119,7 +143,7 @@ class LocalWriter(Writer): """Local Writer Class""" def write_image( - self, name: str, x: TensorType["H", "W", 3], step: int, group: str = None, prefix: str = None + self, name: str, x: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None ) -> None: x = to8b(x) image_path = os.path.join(self.save_dir, f"{name}.jpg") diff --git a/scripts/run_train.py b/scripts/run_train.py index dec47e315..654936a3d 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -16,7 +16,7 @@ import torch.multiprocessing as mp from omegaconf import DictConfig from mattport.nerf.trainer import Trainer -from mattport.utils import profiler +from mattport.utils import comms, profiler logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.DEBUG) @@ -42,12 +42,6 @@ def _set_random_seed(seed) -> None: 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, @@ -94,13 +88,21 @@ def _distributed_worker( logger.error("Process group URL: %s", dist_url) raise e + assert comms._LOCAL_PROCESS_GROUP is None + num_machines = world_size // num_gpus_per_machine + for i in range(num_machines): + ranks_on_i = list(range(i * num_gpus_per_machine, (i + 1) * num_gpus_per_machine)) + pg = dist.new_group(ranks_on_i) + if i == machine_rank: + comms._LOCAL_PROCESS_GROUP = pg + assert num_gpus_per_machine <= torch.cuda.device_count() torch.cuda.set_device(local_rank) _set_random_seed(1234 + local_rank) # NOTE(ethan): will this work for multiple machines? should we use global_rank? - _synchronize(world_size, local_rank) + comms.synchronize(world_size) output = main_func(local_rank, world_size, config) - _synchronize(world_size, local_rank) + comms.synchronize(world_size) dist.destroy_process_group() return output @@ -153,8 +155,7 @@ def launch( except KeyboardInterrupt: pass finally: - if config.logging.enable_profiler: - profiler.PROFILER.print_profile() + profiler.flush_profiler(config.logging) elif world_size > 1: # Using multiple gpus with multiple processes. if dist_url == "auto": @@ -190,8 +191,7 @@ def launch( process.join() logger.info("Process %s finished", str(i)) finally: - if config.logging.enable_profiler: - profiler.PROFILER.print_profile() + profiler.flush_profiler(config.logging) @hydra.main(config_path="../configs", config_name="default.yaml") diff --git a/tests/configs/test_default.yml b/tests/configs/test_default.yml new file mode 100644 index 000000000..3978d3462 --- /dev/null +++ b/tests/configs/test_default.yml @@ -0,0 +1,88 @@ +dry_run: true + +defaults: + # This makes the logging module logs colorful + - override hydra/job_logging: colorlog + - override hydra/hydra_logging: colorlog + +logging_configs: + steps_per_log: 10 + # stats tracker logs configurable training stats per iteration to terminal (e.g. rays/s, PSNR) + enable_stats: True + stats_tracker: + max_history: 10 + stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME] # see mattport/utils/stats_tracker.py for options + # writer logs losses and images per iteration + writer: + type: TensorboardWriter + save_dir: './' + # profiler logs run times of functions and prints at end of training + enable_profiler: True + +hydra: + run: + # Configure output dir of each experiment programmatically from the arguments + # Example "outputs/mnist/classifier/baseline/2021-03-10-141516" + dir: outputs/${experiment_name}/${method_name}/${now:%Y-%m-%d_%H%M%S} + +experiment_name: blender_lego +method_name: vanilla_nerf + +# pretrained loading options +resume_train: + load_dir: # pretrained model directory + load_step: # pretrained model's step number + +max_num_iterations: 10 +steps_per_save: null +steps_per_test: 100 + +model_dir: /tmp/mattport_models/ + +# dataset options +dataset: + data_directory: data/blender/lego_test + dataset_type: blender + downscale_factor: 4 + +# dataloader options +dataloader: + num_workers: 4 + shuffle: true + num_images_to_sample_from: 1 + num_rays_per_batch: 1024 + num_times_to_repeat_images: 40 + +# optimizer options +param_groups: + fields: + optimizer: + type: Adam + lr: 1e-4 + scheduler: + type: StepLR + step_size: 30 # Period of learning rate decay. + gamma: 0.1 # Multiplicative factor of learning rate decay. Default: 0.1. + cameras: + optimizer: + type: SGD + lr: 1e-4 + scheduler: + type: ExponentialLR + gamma: 0.1 + +# machine options +machine_config: + num_gpus: 1 + num_machines: 1 + machine_rank: 0 + dist_url: auto + +# graph options +graph: + # NOTE(ethan): currently setting to nothing b/c we're still designing the config format... + _target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph + near_plane: 1.0 + far_plane: 6.0 + num_coarse_samples: 64 + num_importance_samples: 64