diff --git a/pyrad/engine/trainer.py b/pyrad/engine/trainer.py index 971483178..ee8543ce5 100644 --- a/pyrad/engine/trainer.py +++ b/pyrad/engine/trainer.py @@ -17,7 +17,6 @@ Code to train model. """ import logging import os -from time import time from typing import Callable, Dict, List import torch @@ -34,6 +33,7 @@ from pyrad.utils import profiler, writer from pyrad.utils.callbacks import update_occupancy from pyrad.utils.decorators import check_main_thread from pyrad.utils.misc import instantiate_from_dict_config +from pyrad.utils.writer import EventName, TimeWriter logging.getLogger("PIL").setLevel(logging.WARNING) @@ -179,50 +179,26 @@ class Trainer: return loss_sum def train(self) -> None: - """Main training loop.""" - train_start = time() - num_iterations = self.config.graph.max_num_iterations - iter_dataloader_train = iter(self.dataloader_train) - for step in range(self.start_step, self.start_step + num_iterations): - data_start = time() - ray_indices, batch = next(iter_dataloader_train) - writer.put_time( - name=writer.EventName.ITER_LOAD_TIME, - start_time=data_start, - end_time=time(), - step=step, - avg_over_iters=True, - ) - iter_start = time() - loss_dict = self.train_iteration(ray_indices, batch, step, _callback=[update_occupancy]) - writer.put_time( - name=writer.EventName.RAYS_PER_SEC, - start_time=iter_start, - end_time=time(), - step=step, - avg_over_iters=True, - avg_over_batch=ray_indices.shape[0], - ) - writer.put_time( - name=writer.EventName.ITER_TRAIN_TIME, - start_time=iter_start, - end_time=time(), - step=step, - avg_over_iters=True, - update_eta=True, - ) + """_summary_""" + with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME): + num_iterations = self.config.graph.max_num_iterations + iter_dataloader_train = iter(self.dataloader_train) + for step in range(self.start_step, self.start_step + num_iterations): + with TimeWriter(writer, EventName.ITER_LOAD_TIME, step=step): + ray_indices, batch = next(iter_dataloader_train) - if step != 0 and step % self.config.logging.steps_per_log == 0: - writer.put_dict(name="Loss/train-loss_dict", scalar_dict=loss_dict, step=step) - 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: - self.eval_with_dataloader(self.dataloader_eval, step=step) - self._write_out_storage(step) + with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as t: + loss_dict = self.train_iteration(ray_indices, batch, step, _callback=[update_occupancy]) + writer.put_scalar(name=EventName.RAYS_PER_SEC, scalar=ray_indices.shape[0] / t.duration, step=step) + + if step != 0 and step % self.config.logging.steps_per_log == 0: + writer.put_dict(name="Loss/train-loss_dict", scalar_dict=loss_dict, step=step) + 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: + self.eval_with_dataloader(self.dataloader_eval, step=step) + self._write_out_storage(step) - writer.put_time( - name=writer.EventName.TOTAL_TRAIN_TIME, start_time=train_start, end_time=time(), step=num_iterations - ) self._write_out_storage(num_iterations) def _write_out_storage(self, step: int) -> None: diff --git a/pyrad/utils/writer.py b/pyrad/utils/writer.py index 5c0d36255..ea494dc20 100644 --- a/pyrad/utils/writer.py +++ b/pyrad/utils/writer.py @@ -21,6 +21,7 @@ import enum import os import sys from abc import abstractmethod +from time import time from typing import Dict import imageio @@ -84,42 +85,29 @@ def put_dict(name: str, scalar_dict: float, step: int): @check_main_thread -def put_time( - name: str, - start_time: float, - end_time: float, - step: int, - avg_over_iters: bool = False, - avg_over_batch: int = None, - update_eta: bool = False, -): +def put_time(name: str, duration: float, step: int, avg_over_steps: bool = True, update_eta: bool = False): """Setter function to place a time element into the queue to be written out Processes the time info according to the options: - avg_over_iters (bool): if True, calculate and record a running average of the times - avg_over_batch (int): if set, the size of the batch for which we take the average over (batch/second) + avg_over_steps (bool): if True, calculate and record a running average of the times update_eta (bool): if True, update the ETA. should only be set for the training iterations/s """ if isinstance(name, EventName): name = name.value - GLOBAL_BUFFER["step"] = step - val = end_time - start_time - if avg_over_batch: - val = avg_over_batch / val - - if avg_over_iters: + if avg_over_steps: + GLOBAL_BUFFER["step"] = step curr_event = GLOBAL_BUFFER["events"].get(name, {"buffer": [], "avg": 0}) curr_buffer = curr_event["buffer"] curr_avg = curr_event["avg"] if len(curr_buffer) >= GLOBAL_BUFFER["max_buffer_size"]: curr_buffer.pop(0) - curr_buffer.append(val) + curr_buffer.append(duration) curr_avg = sum(curr_buffer) / len(curr_buffer) put_scalar(name, curr_avg, step) GLOBAL_BUFFER["events"][name] = {"buffer": curr_buffer, "avg": curr_avg} else: - put_scalar(name, val, step) + put_scalar(name, duration, step) if update_eta: ## NOTE: eta should be called with avg train iteration time @@ -204,6 +192,33 @@ class Writer: self.write_scalar(name, scalar, step) +class TimeWriter: + """Timer context manager that calculates duration around wrapped functions""" + + def __init__(self, writer, name, step=None): + self.writer = writer + self.name = name + self.step = step + + self.start = None + self.duration = None + + def __enter__(self): + self.start = time() + return self + + def __exit__(self, *args): + self.duration = time() - self.start + update_step = self.step is not None + self.writer.put_time( + name=self.name, + duration=self.duration, + step=self.step if update_step else GLOBAL_BUFFER["max_iter"], + avg_over_steps=update_step, + update_eta=self.name == EventName.ITER_TRAIN_TIME, + ) + + @decorate_all([check_main_thread]) class WandbWriter(Writer): """WandDB Writer Class"""