From ca76f70d49fe9629a420c7dae36df87efa915713 Mon Sep 17 00:00:00 2001 From: evonneng Date: Fri, 22 Apr 2022 17:33:02 -0700 Subject: [PATCH] Added timing logic within trainer class (#37) * 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 --- mattport/nerf/trainer.py | 56 +++++++++++++++++++++++++++++++++++++--- mattport/utils/writer.py | 24 +---------------- 2 files changed, 53 insertions(+), 27 deletions(-) diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index 4a1c9c5c1..663cb33bd 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -1,8 +1,10 @@ """ Code to train model. """ +import datetime import logging import os +from time import time import torch import torch.distributed as dist @@ -13,10 +15,10 @@ from torch.utils.data import DataLoader from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch from mattport.nerf.dataset.utils import get_dataset_inputs +from mattport.nerf.metrics import get_psnr from mattport.nerf.optimizers import Optimizers from mattport.utils.decorators import check_main_thread from mattport.utils.writer import LocalWriter, TensorboardWriter -from mattport.nerf.metrics import get_psnr logging.getLogger("PIL").setLevel(logging.WARNING) @@ -43,6 +45,7 @@ class Trainer: self.tensorboard_writer = TensorboardWriter( local_rank, world_size, save_dir=os.path.join(os.getcwd(), "writer") ) + self.time_dict = {} def setup_dataset(self): """_summary_""" @@ -121,7 +124,44 @@ class Trainer: ckpt_path, ) + def update_time(self, name, start_time, end_time, num_iter=1, batch_size=None): + """update the time dictionary with running averages/cumulative durations + + Args: + name (_type_): name of function time we are logging + start_time (_type_): start time for the call + end_time (_type_): end time when the call finished executing + num_iter (int, optional): number of total iteration steps. Defaults to 1. + batch_size (_type_, optional): total number of rays in a batch; + if None, reports duration instead of batch per second. Defaults to None. + """ + val = end_time - start_time + if batch_size: + val = batch_size / val + + self.time_dict[name] = (self.time_dict.get(name, 0) * num_iter + val) / (num_iter + 1) + if name == "single iter (time)": + remain_iter = self.config.max_num_iterations - num_iter + self.time_dict["ETA (time)"] = remain_iter * self.time_dict[name] + + def print_times(self): + """helper to print out the timing dictionary.""" + mssg = "" + for k, v in self.time_dict.items(): + if "(time)" in k: + v = str(datetime.timedelta(seconds=v)) + mssg += f"{k}: {v} | " + print(mssg, end="\r") + def get_aggregated_loss(self, losses): + """_summary_ + + Args: + losses (_type_): _description_ + + Returns: + _type_: _description_ + """ loss_sum = 0.0 loss_dict = {} for loss_name in losses.keys(): @@ -132,19 +172,27 @@ class Trainer: def train(self) -> None: """_summary_""" + train_start = time() num_iterations = self.config.max_num_iterations for step in range(self.start_step, self.start_step + num_iterations): + data_start = time() batch = next(iter(self.train_dataloader)) + self.update_time("load data (time)", data_start, time(), num_iter=step) + iter_start = time() loss_dict = self.train_iteration(batch, step) + self.update_time("avg. rays/s (1/s)", iter_start, time(), num_iter=step, batch_size=batch.indices.shape[0]) + self.update_time("single iter (time)", iter_start, time(), num_iter=step) if step != 0 and step % self.config.steps_per_log == 0: self.tensorboard_writer.write_scalar_dict(loss_dict, step, group="Loss", prefix="train-") - logging.info(f"{step}/{num_iterations} with losses {loss_dict}") + logging.info("{%d}/{%d} with losses {%s}", step, num_iterations, loss_dict) # TODO: add the learning rates to tensorboard/logging - if self.config.steps_per_save and step != 0 and step % self.config.steps_per_save == 0: + if step != 0 and self.config.steps_per_save and step % self.config.steps_per_save == 0: self.save_checkpoint(self.config.model_dir, step) - logging.info(f"Saved ckpt at {step}") + logging.info("Saved ckpt at {%d}", step) if step != 0 and step % self.config.steps_per_test == 0: self.test_image(image_idx=0, step=step) + self.print_times() + self.update_time("total train", train_start, time()) def train_iteration(self, batch: dict, step: int): """Run one iteration with a batch of inputs.""" diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index f022dd4df..d11573fc4 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -40,19 +40,6 @@ class Writer: """ raise NotImplementedError - @check_main_thread - def write_text(self, name: str, x: str, step: int, group: str = None, prefix: str = None) -> None: - """Required method to write a string to summary - - Args: - name (str): data identifier - x (str): string 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-" - """ - raise NotImplementedError - @abstractmethod def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None: """Required method to write a single scalar value to the logger @@ -129,15 +116,6 @@ class LocalWriter(Writer): image_path = os.path.join(self.save_dir, f"{name}.jpg") imageio.imwrite(image_path, np.uint8(x.cpu().numpy() * 255.0)) - @check_main_thread - def write_text(self, name: str, x: str, step: int, group: str = None, prefix: str = None) -> None: - """Logs text locally to the terminal in place - Args: - name (str): data identifier - x (str): string to write - """ - print(f"{name}: {x}, step: {step}", end="\r") - @check_main_thread def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None: - self.write_text(name, f"scalar: {scalar:04f}", step) + raise NotImplementedError