Feature/sg 708 time units

src/super_gradients/common/object_names.py

@@ -146,6 +146,7 @@ class Callbacks:
     DETECTION_VISUALIZATION_CALLBACK = "DetectionVisualizationCallback"
     DEKR_VISUALIZATION = "DEKRVisualizationCallback"
     ROBOFLOW_RESULT_CALLBACK = "RoboflowResultCallback"
+    TIMER = "TimerCallback"
 
 
 class LRSchedulers:

src/super_gradients/common/sg_loggers/__init__.py

@@ -3,5 +3,6 @@
 from super_gradients.common.sg_loggers.deci_platform_sg_logger import DeciPlatformSGLogger
 from super_gradients.common.sg_loggers.wandb_sg_logger import WandBSGLogger
 from super_gradients.common.sg_loggers.dagshub_sg_logger import DagsHubSGLogger
+from super_gradients.common.sg_loggers.time_units import TimeUnit, EpochNumber, GlobalBatchStepNumber
 
-__all__ = ["BaseSGLogger", "ClearMLSGLogger", "DeciPlatformSGLogger", "WandBSGLogger", "DagsHubSGLogger"]
+__all__ = ["BaseSGLogger", "ClearMLSGLogger", "DeciPlatformSGLogger", "WandBSGLogger", "DagsHubSGLogger", "TimeUnit", "EpochNumber", "GlobalBatchStepNumber"]

src/super_gradients/common/sg_loggers/abstract_sg_logger.py

@@ -5,6 +5,8 @@
 from PIL import Image
 import torch
 
+from super_gradients.common.sg_loggers.time_units import TimeUnit
+
 
 class AbstractSGLogger(ABC):
     """
@@ -40,7 +42,7 @@ def add_config(self, tag: str, config: dict):
         raise NotImplementedError
 
     @abstractmethod
-    def add_scalar(self, tag: str, scalar_value: float, global_step: int = None):
+    def add_scalar(self, tag: str, scalar_value: float, global_step: Union[int, TimeUnit] = None):
         """
         Add scalar data to SGLogger.
         Typically, this function will add scalar to tensorboard or other experiment management framework.

src/super_gradients/common/sg_loggers/base_sg_logger.py

@@ -10,17 +10,18 @@
 import torch
 from PIL import Image
 
-from super_gradients.common.registry.registry import register_sg_logger
-from super_gradients.common.data_interface.adnn_model_repository_data_interface import ADNNModelRepositoryDataInterfaces
 from super_gradients.common.abstractions.abstract_logger import get_logger
+from super_gradients.common.auto_logging.auto_logger import AutoLoggerConfig
+from super_gradients.common.auto_logging.console_logging import ConsoleSink
+from super_gradients.common.data_interface.adnn_model_repository_data_interface import ADNNModelRepositoryDataInterfaces
 from super_gradients.common.decorators.code_save_decorator import saved_codes
 from super_gradients.common.environment.ddp_utils import multi_process_safe
+from super_gradients.common.environment.monitoring import SystemMonitor
+from super_gradients.common.registry.registry import register_sg_logger
 from super_gradients.common.sg_loggers.abstract_sg_logger import AbstractSGLogger
+from super_gradients.common.sg_loggers.time_units import TimeUnit
 from super_gradients.training.params import TrainingParams
 from super_gradients.training.utils import sg_trainer_utils, get_param
-from super_gradients.common.environment.monitoring import SystemMonitor
-from super_gradients.common.auto_logging.auto_logger import AutoLoggerConfig
-from super_gradients.common.auto_logging.console_logging import ConsoleSink
 
 logger = get_logger(__name__)
 
@@ -155,7 +156,9 @@ def add_config(self, tag: str, config: dict):
         self._write_to_log_file(log_lines)
 
     @multi_process_safe
-    def add_scalar(self, tag: str, scalar_value: float, global_step: int = None):
+    def add_scalar(self, tag: str, scalar_value: float, global_step: Union[int, TimeUnit] = None):
+        if isinstance(global_step, TimeUnit):
+            global_step = global_step.get_value()
         self.tensorboard_writer.add_scalar(tag=tag.lower().replace(" ", "_"), scalar_value=scalar_value, global_step=global_step)
 
     @multi_process_safe

src/super_gradients/common/sg_loggers/clearml_sg_logger.py

@@ -12,6 +12,7 @@
 from super_gradients.common.registry.registry import register_sg_logger
 from super_gradients.common.sg_loggers.base_sg_logger import BaseSGLogger
 from super_gradients.common.environment.ddp_utils import multi_process_safe
+from super_gradients.common.sg_loggers.time_units import TimeUnit
 
 logger = get_logger(__name__)
 
@@ -114,8 +115,10 @@ def __add_scalar(self, tag: str, scalar_value: float, global_step: int):
         self.clearml_logger.report_scalar(title=tag, series=tag, value=scalar_value, iteration=global_step)
 
     @multi_process_safe
-    def add_scalar(self, tag: str, scalar_value: float, global_step: int = 0):
+    def add_scalar(self, tag: str, scalar_value: float, global_step: Union[int, TimeUnit] = 0):
         super(ClearMLSGLogger, self).add_scalar(tag=tag, scalar_value=scalar_value, global_step=global_step)
+        if isinstance(global_step, TimeUnit):
+            global_step = global_step.get_value()
         self.__add_scalar(tag=tag, scalar_value=scalar_value, global_step=global_step)
 
     @multi_process_safe

src/super_gradients/common/sg_loggers/dagshub_sg_logger.py

@@ -9,6 +9,7 @@
 
 from super_gradients.common.sg_loggers.base_sg_logger import BaseSGLogger
 from super_gradients.common.environment.ddp_utils import multi_process_safe
+from super_gradients.common.sg_loggers.time_units import TimeUnit
 
 logger = get_logger(__name__)
 
@@ -177,8 +178,10 @@ def add_config(self, tag: str, config: dict):
                     logger.warning(f"Skip to log {k}: {v}")
 
     @multi_process_safe
-    def add_scalar(self, tag: str, scalar_value: float, global_step: int = 0):
+    def add_scalar(self, tag: str, scalar_value: float, global_step: [int, TimeUnit] = 0):
         super(DagsHubSGLogger, self).add_scalar(tag=tag, scalar_value=scalar_value, global_step=global_step)
+        if isinstance(global_step, TimeUnit):
+            global_step = global_step.get_value()
         mlflow.log_metric(key=tag, value=scalar_value, step=global_step)
 
     @multi_process_safe

src/super_gradients/common/sg_loggers/time_units.py

@@ -0,0 +1,49 @@
+import abc
+import dataclasses
+
+
+class TimeUnit(abc.ABC):
+    """
+    Abstract class for time units. This is used to explicitly log the time unit of a metric/loss.
+    """
+
+    @abc.abstractmethod
+    def get_value(self):
+        ...
+
+    @abc.abstractmethod
+    def get_name(self):
+        ...
+
+
+@dataclasses.dataclass
+class EpochNumber(TimeUnit):
+    """
+    A time unit for epoch number.
+    """
+
+    value: float
+
+    def get_value(self):
+        return self.value
+
+    def get_name(self):
+        return "epoch"
+
+
+@dataclasses.dataclass
+class GlobalBatchStepNumber(TimeUnit):
+    """
+    A time unit for representing total number of batches processed, including training and validation ones.
+    Suppose training loader has 320 batches and validation loader has 80 batches.
+    If the current epoch index is 2 (zero-based), and we are on validation loader and current index is 50 (zero-based),
+    then the global batch step is (320 + 80) * 3 + 320 + 50 = 1570.
+    """
+
+    value: float
+
+    def get_value(self):
+        return self.value
+
+    def get_name(self):
+        return "global_batch_step"

src/super_gradients/common/sg_loggers/wandb_sg_logger.py

@@ -1,18 +1,17 @@
 import os
-
 from typing import Union, Optional, Any
 
-import numpy as np
-from PIL import Image
 import matplotlib.pyplot as plt
+import numpy as np
 import torch
+from PIL import Image
 
-from super_gradients.common.registry.registry import register_sg_logger
-from super_gradients.common.environment.env_variables import env_variables
 from super_gradients.common.abstractions.abstract_logger import get_logger
-
-from super_gradients.common.sg_loggers.base_sg_logger import BaseSGLogger
 from super_gradients.common.environment.ddp_utils import multi_process_safe
+from super_gradients.common.environment.env_variables import env_variables
+from super_gradients.common.registry.registry import register_sg_logger
+from super_gradients.common.sg_loggers.base_sg_logger import BaseSGLogger
+from super_gradients.common.sg_loggers.time_units import TimeUnit
 
 logger = get_logger(__name__)
 
@@ -169,9 +168,12 @@ def add_config(self, tag: str, config: dict):
         wandb.config.update(config, allow_val_change=self.resumed)
 
     @multi_process_safe
-    def add_scalar(self, tag: str, scalar_value: float, global_step: int = 0):
+    def add_scalar(self, tag: str, scalar_value: float, global_step: Union[int, TimeUnit] = 0):
         super(WandBSGLogger, self).add_scalar(tag=tag, scalar_value=scalar_value, global_step=global_step)
-        wandb.log(data={tag: scalar_value}, step=global_step)
+        if isinstance(global_step, TimeUnit):
+            wandb.log(data={tag: scalar_value, global_step.get_name(): global_step.get_value()})
+        else:
+            wandb.log(data={tag: scalar_value}, step=global_step)
 
     @multi_process_safe
     def add_scalars(self, tag_scalar_dict: dict, global_step: int = 0):

src/super_gradients/training/sg_trainer/sg_trainer.py

@@ -438,6 +438,7 @@ def _train_epoch(self, epoch: int, silent_mode: bool = False) -> tuple:
                 lr_warmup_epochs=self.training_params.lr_warmup_epochs,
                 sg_logger=self.sg_logger,
                 train_loader=self.train_loader,
+                valid_loader=self.valid_loader,
                 context_methods=self._get_context_methods(Phase.TRAIN_BATCH_END),
                 ddp_silent_mode=self.ddp_silent_mode,
             )
@@ -1835,6 +1836,8 @@ def evaluate(
             device=device_config.device,
             lr_warmup_epochs=lr_warmup_epochs,
             sg_logger=self.sg_logger,
+            train_loader=self.train_loader,
+            valid_loader=self.valid_loader,
             context_methods=self._get_context_methods(Phase.VALIDATION_BATCH_END),
         )
 

src/super_gradients/training/utils/callbacks/__init__.py

@@ -21,6 +21,7 @@
     TrainingStageSwitchCallbackBase,
     YoloXTrainingStageSwitchCallback,
     TestLRCallback,
+    TimerCallback,
 )
 from super_gradients.training.utils.callbacks.ppyoloe_switch_callback import PPYoloETrainingStageSwitchCallback
 from super_gradients.common.object_names import Callbacks, LRSchedulers, LRWarmups
@@ -60,4 +61,5 @@
     "CallbackHandler",
     "TestLRCallback",
     "PPYoloETrainingStageSwitchCallback",
+    "TimerCallback",
 ]

src/super_gradients/training/utils/callbacks/callbacks.py

@@ -17,6 +17,7 @@
 from super_gradients.common.plugins.deci_client import DeciClient
 from super_gradients.common.registry.registry import register_lr_scheduler, register_lr_warmup, register_callback
 from super_gradients.common.object_names import LRSchedulers, LRWarmups, Callbacks
+from super_gradients.common.sg_loggers.time_units import GlobalBatchStepNumber, EpochNumber
 from super_gradients.training.utils.callbacks.base_callbacks import PhaseCallback, PhaseContext, Phase, Callback
 from super_gradients.training.utils.detection_utils import DetectionVisualization, DetectionPostPredictionCallback
 from super_gradients.training.utils.segmentation_utils import BinarySegmentationVisualization
@@ -759,3 +760,102 @@ def __init__(self, lr_placeholder):
 
     def __call__(self, context: PhaseContext):
         self.lr_placeholder.append(context.optimizer.param_groups[0]["lr"])
+
+
+@register_callback(Callbacks.TIMER)
+class TimerCallback(Callback):
+    def __init__(self):
+        self.events = {}
+
+    @multi_process_safe
+    def on_train_loader_start(self, context: PhaseContext) -> None:
+        self.events["on_train_loader_start"] = cv2.getTickCount()
+
+    @multi_process_safe
+    def on_train_batch_start(self, context: PhaseContext) -> None:
+        self.events["on_train_batch_start"] = cv2.getTickCount()
+
+    @multi_process_safe
+    def on_train_batch_loss_end(self, context: PhaseContext) -> None:
+        self.events["on_train_batch_loss_end"] = cv2.getTickCount()
+        context.sg_logger.add_scalar(
+            tag="timer/train_batch_forward_with_loss_ms",
+            scalar_value=self._elapsed_time_between("on_train_batch_start", "on_train_batch_loss_end"),
+            global_step=GlobalBatchStepNumber(self._infer_global_step(context, is_train_loader=True)),
+        )
+
+    @multi_process_safe
+    def on_train_batch_gradient_step_start(self, context: PhaseContext) -> None:
+        self.events["on_train_batch_gradient_step_start"] = cv2.getTickCount()
+
+    @multi_process_safe
+    def on_train_batch_gradient_step_end(self, context: PhaseContext) -> None:
+        self.events["on_train_batch_gradient_step_end"] = cv2.getTickCount()
+        context.sg_logger.add_scalar(
+            tag="timer/train_batch_gradient_time",
+            scalar_value=self._elapsed_time_between("on_train_batch_gradient_step_start", "on_train_batch_gradient_step_end"),
+            global_step=GlobalBatchStepNumber(self._infer_global_step(context, is_train_loader=True)),
+        )
+
+    @multi_process_safe
+    def on_train_batch_end(self, context: PhaseContext) -> None:
+        self.events["on_train_batch_end"] = cv2.getTickCount()
+        context.sg_logger.add_scalar(
+            tag="timer/train_batch_total_time_ms",
+            scalar_value=self._elapsed_time_between("on_train_batch_start", "on_train_batch_end"),
+            global_step=GlobalBatchStepNumber(self._infer_global_step(context, is_train_loader=True)),
+        )
+
+    @multi_process_safe
+    def on_train_loader_end(self, context: PhaseContext) -> None:
+        self.events["on_train_loader_end"] = cv2.getTickCount()
+        context.sg_logger.add_scalar(
+            tag="timer/train_loader_total_time_ms",
+            scalar_value=self._elapsed_time_between("on_train_loader_start", "on_train_loader_end"),
+            global_step=EpochNumber(context.epoch),
+        )
+
+    @multi_process_safe
+    def on_validation_loader_start(self, context: PhaseContext) -> None:
+        self.events["on_validation_loader_start"] = cv2.getTickCount()
+
+    @multi_process_safe
+    def on_validation_batch_start(self, context: PhaseContext) -> None:
+        self.events["on_validation_batch_start"] = cv2.getTickCount()
+
+    @multi_process_safe
+    def on_validation_batch_end(self, context: PhaseContext) -> None:
+        self.events["on_validation_batch_end"] = cv2.getTickCount()
+        context.sg_logger.add_scalar(
+            tag="timer/validation_batch_total_time_ms",
+            scalar_value=self._elapsed_time_between("on_validation_batch_start", "on_validation_batch_end"),
+            global_step=GlobalBatchStepNumber(self._infer_global_step(context, is_train_loader=False)),
+        )
+
+    @multi_process_safe
+    def on_validation_loader_end(self, context: PhaseContext) -> None:
+        self.events["on_validation_loader_end"] = cv2.getTickCount()
+        context.sg_logger.add_scalar(
+            tag="timer/validation_loader_total_time_ms",
+            scalar_value=self._elapsed_time_between("on_validation_loader_start", "on_validation_loader_end"),
+            global_step=EpochNumber(context.epoch),
+        )
+
+        context.sg_logger.add_scalar(
+            tag="timer/epoch_total_time_sec",
+            scalar_value=self._elapsed_time_between("on_train_loader_start", "on_validation_loader_end") / 1000.0,
+            global_step=EpochNumber(context.epoch),
+        )
+
+    def _elapsed_time_between(self, start_event, end_event):
+        return 1000.0 * (self.events[end_event] - self.events[start_event]) / cv2.getTickFrequency()
+
+    def _infer_global_step(self, context: PhaseContext, is_train_loader: bool):
+        train_loader_length = len(context.train_loader) if context.train_loader is not None else 0
+        valid_loader_length = len(context.valid_loader) if context.valid_loader is not None else 0
+        total_steps_in_epoch = train_loader_length + valid_loader_length
+        total_steps_in_done = context.epoch * total_steps_in_epoch
+        if is_train_loader:
+            return total_steps_in_done + context.batch_idx
+        else:
+            return total_steps_in_done + train_loader_length + context.batch_idx