Cleaner timing function calls with wrapped logic (#39)

* cleanup timing for logger events

* cleaned up context manager timer

* fixed lint

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:

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"""