Create and use a context manager to time code

Use the timer context manager in all places where code was being timed

- Benefits
  - Deduplicate timing code scattered across codebase.
  - Provides single place to manage perf timing code
  - Use consistent timing log patterns

src/utils/helpers.py

@@ -2,10 +2,12 @@
 from __future__ import annotations  # to avoid quoting type hints
 import logging
 import sys
+import torch
 from collections import OrderedDict
 from importlib import import_module
 from os.path import join
 from pathlib import Path
+from time import perf_counter
 from typing import Optional, Union, TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -81,6 +83,25 @@ def get_class_by_name(name: str) -> object:
     return getattr(import_module(module_name), class_name)
 
 
+class timer:
+    '''Context manager to log time taken for a block of code to run'''
+    def __init__(self, message: str, logger: logging.Logger, device: torch.device = None):
+        self.message = message
+        self.logger = logger
+        self.device = device
+
+    def __enter__(self):
+        self.start = perf_counter()
+        return self
+
+    def __exit__(self, *_):
+        elapsed = perf_counter() - self.start
+        if self.device is None:
+            self.logger.debug(f"{self.message}: {elapsed:.3f} seconds")
+        else:
+            self.logger.debug(f"{self.message}: {elapsed:.3f} seconds on device: {self.device}")
+
+
 class LRU(OrderedDict):
     def __init__(self, *args, capacity=128, **kwargs):
         self.capacity = capacity