def get_builtin_activation_type(activation: Union[str, None], **kwargs) -> Type[nn.Module]:
    Returns activation class by its name from torch.nn namespace. This function support all modules available from
    torch.nn and also their lower-case aliases.
    On top of that, it supports a few aliaes: leaky_relu (LeakyReLU), swish (silu).
    >>> act_cls = get_activation_type("LeakyReLU", inplace=True, slope=0.01)
    >>> act = act_cls()
    :param activation: Activation function name (E.g. ReLU). If None - return nn.Identity
    :param **kwargs  : Extra arguments to pass to constructor during instantiation (E.g. inplace=True)
    :returns         : Type of the activation function that is ready to be instantiated
    if activation is None:
        activation_cls = nn.Identity
    else:
        lowercase_aliases: Dict[str, str] = dict((k.lower(), k) for k in torch.nn.__dict__.keys())
        # Register additional aliases
        lowercase_aliases["leaky_relu"] = "LeakyReLU"  # LeakyRelu in snake_case
        lowercase_aliases["swish"] = "SiLU"  # Swish shich is equivalent to SiLU
        lowercase_aliases["none"] = "Identity"
        if activation in lowercase_aliases:
            activation = lowercase_aliases[activation]
        if activation not in torch.nn.__dict__:
            raise KeyError(f"Requested activation function {activation} is not known")
        activation_cls = torch.nn.__dict__[activation]
        if len(kwargs):
            activation_cls = partial(activation_cls, **kwargs)
    return activation_cls
29

def batch_distance2bbox(points: Tensor, distance: Tensor, max_shapes: Optional[Tensor] = None) -> Tensor:
    """Decode distance prediction to bounding box for batch.
    :param points: [B, ..., 2], "xy" format
    :param distance: [B, ..., 4], "ltrb" format
    :param max_shapes: [B, 2], "h,w" format, Shape of the image.
    :return: Tensor: Decoded bboxes, "x1y1x2y2" format.
    lt, rb = torch.split(distance, 2, dim=-1)
    # while tensor add parameters, parameters should be better placed on the second place
    x1y1 = -lt + points
    x2y2 = rb + points
    out_bbox = torch.cat([x1y1, x2y2], dim=-1)
    if max_shapes is not None:
        max_shapes = max_shapes.flip(-1).tile([1, 2])
        delta_dim = out_bbox.ndim - max_shapes.ndim
        for _ in range(delta_dim):
            max_shapes.unsqueeze_(1)
        out_bbox = torch.where(out_bbox < max_shapes, out_bbox, max_shapes)
        out_bbox = torch.where(out_bbox > 0, out_bbox, torch.zeros_like(out_bbox))
    return out_bbox
      
Base callback class with all the callback methods. Derived classes may override one or many of the available events
to receive callbacks when such events are triggered by the training loop.
The order of the events is as follows:
on_training_start(context)                              # called once before training starts, good for setting up the warmup LR
for epoch in range(epochs):
    on_train_loader_start(context)
        for batch in train_loader:
            on_train_batch_start(context)
            on_train_batch_loss_end(context)               # called after loss has been computed
            on_train_batch_backward_end(context)           # called after .backward() was called
            on_train_batch_gradient_step_start(context)    # called before the optimizer step about to happen (gradient clipping, logging of gradients)
            on_train_batch_gradient_step_end(context)      # called after gradient step was done, good place to update LR (for step-based schedulers)
            on_train_batch_end(context)
    on_train_loader_end(context)
    on_validation_loader_start(context)
        for batch in validation_loader:
            on_validation_batch_start(context)
            on_validation_batch_end(context)
    on_validation_loader_end(context)
    on_validation_end_best_epoch(context)
on_test_start(context)
    for batch in test_loader:
        on_test_batch_start(context)
        on_test_batch_end(context)
on_test_end(context)
on_training_end(context)                    # called once after training ends.
Correspondence mapping from the old callback API:
on_training_start(context)                                 <-> Phase.PRE_TRAINING
for epoch in range(epochs):
    on_train_loader_start(context)                         <-> Phase.TRAIN_EPOCH_START
        for batch in train_loader:
            on_train_batch_start(context)
            on_train_batch_loss_end(context)
            on_train_batch_backward_end(context)           <-> Phase.TRAIN_BATCH_END
            on_train_batch_gradient_step_start(context)
            on_train_batch_gradient_step_end(context)      <-> Phase.TRAIN_BATCH_STEP
            on_train_batch_end(context)
    on_train_loader_end(context)                           <-> Phase.TRAIN_EPOCH_END
on_validation_loader_start(context)
    for batch in validation_loader:
        on_validation_batch_start(context)
        on_validation_batch_end(context)               <-> Phase.VALIDATION_BATCH_END
on_validation_loader_end(context)                      <-> Phase.VALIDATION_EPOCH_END
on_validation_end_best_epoch(context)                  <-> Phase.VALIDATION_END_BEST_EPOCH
on_test_start(context)
    for batch in test_loader:
        on_test_batch_start(context)
        on_test_batch_end(context)                         <-> Phase.TEST_BATCH_END
on_test_end(context)                                       <-> Phase.TEST_END
on_training_end(context)                                   <-> Phase.POST_TRAINING
          
Source code in V3_2/src/super_gradients/training/utils/callbacks/base_callbacks.py
355

class Callback:
    Base callback class with all the callback methods. Derived classes may override one or many of the available events
    to receive callbacks when such events are triggered by the training loop.
    The order of the events is as follows:
    on_training_start(context)                              # called once before training starts, good for setting up the warmup LR
        for epoch in range(epochs):
            on_train_loader_start(context)
                for batch in train_loader:
                    on_train_batch_start(context)
                    on_train_batch_loss_end(context)               # called after loss has been computed
                    on_train_batch_backward_end(context)           # called after .backward() was called
                    on_train_batch_gradient_step_start(context)    # called before the optimizer step about to happen (gradient clipping, logging of gradients)
                    on_train_batch_gradient_step_end(context)      # called after gradient step was done, good place to update LR (for step-based schedulers)
                    on_train_batch_end(context)
            on_train_loader_end(context)
            on_validation_loader_start(context)
                for batch in validation_loader:
                    on_validation_batch_start(context)
                    on_validation_batch_end(context)
            on_validation_loader_end(context)
            on_validation_end_best_epoch(context)
        on_test_start(context)
            for batch in test_loader:
                on_test_batch_start(context)
                on_test_batch_end(context)
        on_test_end(context)
    on_training_end(context)                    # called once after training ends.
    Correspondence mapping from the old callback API:
    on_training_start(context)                                 <-> Phase.PRE_TRAINING
    for epoch in range(epochs):
        on_train_loader_start(context)                         <-> Phase.TRAIN_EPOCH_START
            for batch in train_loader:
                on_train_batch_start(context)
                on_train_batch_loss_end(context)
                on_train_batch_backward_end(context)           <-> Phase.TRAIN_BATCH_END
                on_train_batch_gradient_step_start(context)
                on_train_batch_gradient_step_end(context)      <-> Phase.TRAIN_BATCH_STEP
                on_train_batch_end(context)
        on_train_loader_end(context)                           <-> Phase.TRAIN_EPOCH_END
        on_validation_loader_start(context)
            for batch in validation_loader:
                on_validation_batch_start(context)
                on_validation_batch_end(context)               <-> Phase.VALIDATION_BATCH_END
        on_validation_loader_end(context)                      <-> Phase.VALIDATION_EPOCH_END
        on_validation_end_best_epoch(context)                  <-> Phase.VALIDATION_END_BEST_EPOCH
    on_test_start(context)
        for batch in test_loader:
            on_test_batch_start(context)
            on_test_batch_end(context)                         <-> Phase.TEST_BATCH_END
    on_test_end(context)                                       <-> Phase.TEST_END
    on_training_end(context)                                   <-> Phase.POST_TRAINING
    def on_training_start(self, context: PhaseContext) -> None:
        Called once before start of the first epoch
        At this point, the context argument is guaranteed to have the following attributes:
        - optimizer
        - net
        - checkpoints_dir_path
        - criterion
        - sg_logger
        - train_loader
        - valid_loader
        - training_params
        - checkpoint_params
        - architecture
        - arch_params
        - metric_to_watch
        - device
        - ema_model
        The corresponding Phase enum value for this event is Phase.PRE_TRAINING.
        :param context:
        :return:
    def on_train_loader_start(self, context: PhaseContext) -> None:
        Called each epoch at the start of train data loader (before getting the first batch).
        At this point, the context argument is guaranteed to have the following attributes:
        - epoch
        The corresponding Phase enum value for this event is Phase.TRAIN_EPOCH_START.
        :param context:
        :return:
    def on_train_batch_start(self, context: PhaseContext) -> None:
        Called at each batch after getting batch of data from data loader and moving it to target device.
        This event triggered AFTER Trainer.pre_prediction_callback call (If it was defined).
        At this point the context argument is guaranteed to have the following attributes:
        - batch_idx
        - inputs
        - targets
        - **additional_batch_items
        :param context:
        :return:
    def on_train_batch_loss_end(self, context: PhaseContext) -> None:
        Called after model forward and loss computation has been done.
        At this point the context argument is guaranteed to have the following attributes:
        - preds
        - loss_log_items
        The corresponding Phase enum value for this event is Phase.TRAIN_BATCH_END.
        :param context:
        :return:
    def on_train_batch_backward_end(self, context: PhaseContext) -> None:
        Called after loss.backward() method was called for a given batch
        :param context:
        :return:
    def on_train_batch_gradient_step_start(self, context: PhaseContext) -> None:
        Called before the graadient step is about to happen.
        Good place to clip gradients (with respect to scaler), log gradients to data ratio, etc.
        :param context:
        :return:
    def on_train_batch_gradient_step_end(self, context: PhaseContext) -> None:
        Called after gradient step has been performed. Good place to update LR (for step-based schedulers)
        The corresponding Phase enum value for this event is Phase.TRAIN_BATCH_STEP.
        :param context:
        :return:
    def on_train_batch_end(self, context: PhaseContext) -> None:
        Called after all forward/backward/optimizer steps have been performed for a given batch and there is nothing left to do.
        :param context:
        :return:
    def on_train_loader_end(self, context: PhaseContext) -> None:
        Called each epoch at the end of train data loader (after processing the last batch).
        The corresponding Phase enum value for this event is Phase.TRAIN_EPOCH_END.
        :param context:
        :return:
    def on_validation_loader_start(self, context: PhaseContext) -> None:
        Called each epoch at the start of validation data loader (before getting the first batch).
        :param context:
        :return:
    def on_validation_batch_start(self, context: PhaseContext) -> None:
        Called at each batch after getting batch of data from validation loader and moving it to target device.
        :param context:
        :return:
    def on_validation_batch_end(self, context: PhaseContext) -> None:
        Called after all forward step / loss / metric computation have been performed for a given batch and there is nothing left to do.
        The corresponding Phase enum value for this event is Phase.VALIDATION_BATCH_END.
        :param context:
        :return:
    def on_validation_loader_end(self, context: PhaseContext) -> None:
        Called each epoch at the end of validation data loader (after processing the last batch).
        The corresponding Phase enum value for this event is Phase.VALIDATION_EPOCH_END.
        :param context:
        :return:
    def on_validation_end_best_epoch(self, context: PhaseContext) -> None:
        Called each epoch after validation has been performed and the best metric has been achieved.
        The corresponding Phase enum value for this event is Phase.VALIDATION_END_BEST_EPOCH.
        :param context:
        :return:
    def on_test_loader_start(self, context: PhaseContext) -> None:
        Called once at the start of test data loader (before getting the first batch).
        :param context:
        :return:
    def on_test_batch_start(self, context: PhaseContext) -> None:
        Called at each batch after getting batch of data from test loader and moving it to target device.
        :param context:
        :return:
    def on_test_batch_end(self, context: PhaseContext) -> None:
        Called after all forward step have been performed for a given batch and there is nothing left to do.
        The corresponding Phase enum value for this event is Phase.TEST_BATCH_END.
        :param context:
        :return:
    def on_test_loader_end(self, context: PhaseContext) -> None:
        Called once at the end of test data loader (after processing the last batch).
        The corresponding Phase enum value for this event is Phase.TEST_END.
        :param context:
        :return:
    def on_training_end(self, context: PhaseContext) -> None:
        Called once after the training loop has finished (Due to reaching optimization criterion or because of an error.)
        The corresponding Phase enum value for this event is Phase.POST_TRAINING.
        :param context:
        :return:
      
Called after all forward step have been performed for a given batch and there is nothing left to do.
The corresponding Phase enum value for this event is Phase.TEST_BATCH_END.
  
Parameters:
        Description
        Default
337

def on_test_batch_end(self, context: PhaseContext) -> None:
    Called after all forward step have been performed for a given batch and there is nothing left to do.
    The corresponding Phase enum value for this event is Phase.TEST_BATCH_END.
    :param context:
    :return:
      
Called at each batch after getting batch of data from test loader and moving it to target device.
  
Parameters:
        Description
        Default
328

def on_test_batch_start(self, context: PhaseContext) -> None:
    Called at each batch after getting batch of data from test loader and moving it to target device.
    :param context:
    :return:
      
Called once at the end of test data loader (after processing the last batch).
The corresponding Phase enum value for this event is Phase.TEST_END.
  
Parameters:
        Description
        Default
346

def on_test_loader_end(self, context: PhaseContext) -> None:
    Called once at the end of test data loader (after processing the last batch).
    The corresponding Phase enum value for this event is Phase.TEST_END.
    :param context:
    :return:
320

def on_test_loader_start(self, context: PhaseContext) -> None:
    Called once at the start of test data loader (before getting the first batch).
    :param context:
    :return:
229

def on_train_batch_backward_end(self, context: PhaseContext) -> None:
    Called after loss.backward() method was called for a given batch
    :param context:
    :return:
      
Called after all forward/backward/optimizer steps have been performed for a given batch and there is nothing left to do.
  
Parameters:
        Description
        Default
257

def on_train_batch_end(self, context: PhaseContext) -> None:
    Called after all forward/backward/optimizer steps have been performed for a given batch and there is nothing left to do.
    :param context:
    :return:
      
Called after gradient step has been performed. Good place to update LR (for step-based schedulers)
The corresponding Phase enum value for this event is Phase.TRAIN_BATCH_STEP.
  
Parameters:
        Description
        Default
247

def on_train_batch_gradient_step_end(self, context: PhaseContext) -> None:
    Called after gradient step has been performed. Good place to update LR (for step-based schedulers)
    The corresponding Phase enum value for this event is Phase.TRAIN_BATCH_STEP.
    :param context:
    :return:
      
Called before the graadient step is about to happen.
Good place to clip gradients (with respect to scaler), log gradients to data ratio, etc.
  
Parameters:
        Description
        Default
238

def on_train_batch_gradient_step_start(self, context: PhaseContext) -> None:
    Called before the graadient step is about to happen.
    Good place to clip gradients (with respect to scaler), log gradients to data ratio, etc.
    :param context:
    :return:
      
Called after model forward and loss computation has been done.
At this point the context argument is guaranteed to have the following attributes:
- preds
- loss_log_items
The corresponding Phase enum value for this event is Phase.TRAIN_BATCH_END.
  
Parameters:
        Description
        Default
220

def on_train_batch_loss_end(self, context: PhaseContext) -> None:
    Called after model forward and loss computation has been done.
    At this point the context argument is guaranteed to have the following attributes:
    - preds
    - loss_log_items
    The corresponding Phase enum value for this event is Phase.TRAIN_BATCH_END.
    :param context:
    :return:
      
Called at each batch after getting batch of data from data loader and moving it to target device.
This event triggered AFTER Trainer.pre_prediction_callback call (If it was defined).
At this point the context argument is guaranteed to have the following attributes:
- batch_idx
- inputs
- targets
- **additional_batch_items
  
Parameters:
        Description
        Default
206

def on_train_batch_start(self, context: PhaseContext) -> None:
    Called at each batch after getting batch of data from data loader and moving it to target device.
    This event triggered AFTER Trainer.pre_prediction_callback call (If it was defined).
    At this point the context argument is guaranteed to have the following attributes:
    - batch_idx
    - inputs
    - targets
    - **additional_batch_items
    :param context:
    :return:
      
Called each epoch at the end of train data loader (after processing the last batch).
The corresponding Phase enum value for this event is Phase.TRAIN_EPOCH_END.
  
Parameters:
        Description
        Default
267

def on_train_loader_end(self, context: PhaseContext) -> None:
    Called each epoch at the end of train data loader (after processing the last batch).
    The corresponding Phase enum value for this event is Phase.TRAIN_EPOCH_END.
    :param context:
    :return:
      
Called each epoch at the start of train data loader (before getting the first batch).
At this point, the context argument is guaranteed to have the following attributes:
- epoch
The corresponding Phase enum value for this event is Phase.TRAIN_EPOCH_START.
  
Parameters:
        Description
        Default
190

def on_train_loader_start(self, context: PhaseContext) -> None:
    Called each epoch at the start of train data loader (before getting the first batch).
    At this point, the context argument is guaranteed to have the following attributes:
    - epoch
    The corresponding Phase enum value for this event is Phase.TRAIN_EPOCH_START.
    :param context:
    :return:
      
Called once after the training loop has finished (Due to reaching optimization criterion or because of an error.)
The corresponding Phase enum value for this event is Phase.POST_TRAINING.
  
Parameters:
        Description
        Default
355

def on_training_end(self, context: PhaseContext) -> None:
    Called once after the training loop has finished (Due to reaching optimization criterion or because of an error.)
    The corresponding Phase enum value for this event is Phase.POST_TRAINING.
    :param context:
    :return:
      
Called once before start of the first epoch
At this point, the context argument is guaranteed to have the following attributes:
- optimizer
- net
- checkpoints_dir_path
- criterion
- sg_logger
- train_loader
- valid_loader
- training_params
- checkpoint_params
- architecture
- arch_params
- metric_to_watch
- device
- ema_model
The corresponding Phase enum value for this event is Phase.PRE_TRAINING.
  
Parameters:
        Description
        Default
179

def on_training_start(self, context: PhaseContext) -> None:
    Called once before start of the first epoch
    At this point, the context argument is guaranteed to have the following attributes:
    - optimizer
    - net
    - checkpoints_dir_path
    - criterion
    - sg_logger
    - train_loader
    - valid_loader
    - training_params
    - checkpoint_params
    - architecture
    - arch_params
    - metric_to_watch
    - device
    - ema_model
    The corresponding Phase enum value for this event is Phase.PRE_TRAINING.
    :param context:
    :return:
      
Called after all forward step / loss / metric computation have been performed for a given batch and there is nothing left to do.
The corresponding Phase enum value for this event is Phase.VALIDATION_BATCH_END.
  
Parameters:
        Description
        Default
293

def on_validation_batch_end(self, context: PhaseContext) -> None:
    Called after all forward step / loss / metric computation have been performed for a given batch and there is nothing left to do.
    The corresponding Phase enum value for this event is Phase.VALIDATION_BATCH_END.
    :param context:
    :return:
      
Called at each batch after getting batch of data from validation loader and moving it to target device.
  
Parameters:
        Description
        Default
284

def on_validation_batch_start(self, context: PhaseContext) -> None:
    Called at each batch after getting batch of data from validation loader and moving it to target device.
    :param context:
    :return:
      
Called each epoch after validation has been performed and the best metric has been achieved.
The corresponding Phase enum value for this event is Phase.VALIDATION_END_BEST_EPOCH.
  
Parameters:
        Description
        Default
311

def on_validation_end_best_epoch(self, context: PhaseContext) -> None:
    Called each epoch after validation has been performed and the best metric has been achieved.
    The corresponding Phase enum value for this event is Phase.VALIDATION_END_BEST_EPOCH.
    :param context:
    :return:
      
Called each epoch at the end of validation data loader (after processing the last batch).
The corresponding Phase enum value for this event is Phase.VALIDATION_EPOCH_END.
  
Parameters:
        Description
        Default
302

def on_validation_loader_end(self, context: PhaseContext) -> None:
    Called each epoch at the end of validation data loader (after processing the last batch).
    The corresponding Phase enum value for this event is Phase.VALIDATION_EPOCH_END.
    :param context:
    :return:
      
Called each epoch at the start of validation data loader (before getting the first batch).
  
Parameters:
        Description
        Default
276

def on_validation_loader_start(self, context: PhaseContext) -> None:
    Called each epoch at the start of validation data loader (before getting the first batch).
    :param context:
    :return:
    def __init__(self, callbacks: List[Callback]):
        # TODO: Add reordering of callbacks to make sure that they are called in the right order
        # For instance, two callbacks may be dependent on each other, so the first one should be called first
        # Example: Gradient Clipping & Gradient Logging callback. We first need to clip the gradients, and then log them
        # So if user added them in wrong order we can guarantee their order would be correct.
        # We can achieve this by adding a property to the callback to the callback indicating it's priority:
        # Forward   = 0
        # Loss      = 100
        # Backward  = 200
        # Metrics   = 300
        # Scheduler = 400
        # Logging   = 500
        # So ordering callbacks by their order would ensure than we first run all Forward-related callbacks (for a given event),
        # Than backward, and only then - logging.
        self.callbacks = callbacks
    def on_training_start(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_training_start(context)
    def on_train_loader_start(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_train_loader_start(context)
    def on_train_batch_start(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_train_batch_start(context)
    def on_train_batch_loss_end(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_train_batch_loss_end(context)
    def on_train_batch_backward_end(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_train_batch_backward_end(context)
    def on_train_batch_gradient_step_start(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_train_batch_gradient_step_start(context)
    def on_train_batch_gradient_step_end(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_train_batch_gradient_step_end(context)
    def on_train_batch_end(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_train_batch_end(context)
    def on_validation_loader_start(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_validation_loader_start(context)
    def on_validation_batch_start(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_validation_batch_start(context)
    def on_validation_batch_end(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_validation_batch_end(context)
    def on_validation_loader_end(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_validation_loader_end(context)
    def on_train_loader_end(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_train_loader_end(context)
    def on_training_end(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_training_end(context)
    def on_validation_end_best_epoch(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_validation_end_best_epoch(context)
    def on_test_loader_start(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_test_loader_start(context)
    def on_test_batch_start(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_test_batch_start(context)
    def on_test_batch_end(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_test_batch_end(context)
    def on_test_loader_end(self, context: PhaseContext) -> None:
        for callback in self.callbacks:
            callback.on_test_loader_end(context)
      
Kept here to keep backward compatibility with old code. New callbacks should use Callback class instead.
This callback supports receiving only a subset of events defined in Phase enum:
PRE_TRAINING = "PRE_TRAINING"
TRAIN_EPOCH_START = "TRAIN_EPOCH_START"
TRAIN_BATCH_END = "TRAIN_BATCH_END"
TRAIN_BATCH_STEP = "TRAIN_BATCH_STEP"
TRAIN_EPOCH_END = "TRAIN_EPOCH_END"
VALIDATION_BATCH_END = "VALIDATION_BATCH_END"
VALIDATION_EPOCH_END = "VALIDATION_EPOCH_END"
VALIDATION_END_BEST_EPOCH = "VALIDATION_END_BEST_EPOCH"
TEST_BATCH_END = "TEST_BATCH_END"
TEST_END = "TEST_END"
POST_TRAINING = "POST_TRAINING"
          
Source code in V3_2/src/super_gradients/training/utils/callbacks/base_callbacks.py
429

class PhaseCallback(Callback):
    Kept here to keep backward compatibility with old code. New callbacks should use Callback class instead.
    This callback supports receiving only a subset of events defined in Phase enum:
    PRE_TRAINING = "PRE_TRAINING"
    TRAIN_EPOCH_START = "TRAIN_EPOCH_START"
    TRAIN_BATCH_END = "TRAIN_BATCH_END"
    TRAIN_BATCH_STEP = "TRAIN_BATCH_STEP"
    TRAIN_EPOCH_END = "TRAIN_EPOCH_END"
    VALIDATION_BATCH_END = "VALIDATION_BATCH_END"
    VALIDATION_EPOCH_END = "VALIDATION_EPOCH_END"
    VALIDATION_END_BEST_EPOCH = "VALIDATION_END_BEST_EPOCH"
    TEST_BATCH_END = "TEST_BATCH_END"
    TEST_END = "TEST_END"
    POST_TRAINING = "POST_TRAINING"
    def __init__(self, phase: Phase):
        self.phase = phase
    def __call__(self, *args, **kwargs):
        raise NotImplementedError
    def __repr__(self) -> str:
        return self.__class__.__name__
    def on_training_start(self, context: PhaseContext) -> None:
        if self.phase == Phase.PRE_TRAINING:
            self(context)
    def on_train_loader_start(self, context: PhaseContext) -> None:
        if self.phase == Phase.TRAIN_EPOCH_START:
            self(context)
    def on_train_batch_loss_end(self, context: PhaseContext) -> None:
        if self.phase == Phase.TRAIN_BATCH_END:
            self(context)
    def on_train_batch_gradient_step_end(self, context: PhaseContext) -> None:
        if self.phase == Phase.TRAIN_BATCH_STEP:
            self(context)
    def on_train_loader_end(self, context: PhaseContext) -> None:
        if self.phase == Phase.TRAIN_EPOCH_END:
            self(context)
    def on_validation_batch_end(self, context: PhaseContext) -> None:
        if self.phase == Phase.VALIDATION_BATCH_END:
            self(context)
    def on_validation_loader_end(self, context: PhaseContext) -> None:
        if self.phase == Phase.VALIDATION_EPOCH_END:
            self(context)
    def on_validation_end_best_epoch(self, context: PhaseContext) -> None:
        if self.phase == Phase.VALIDATION_END_BEST_EPOCH:
            self(context)
    def on_test_batch_end(self, context: PhaseContext) -> None:
        if self.phase == Phase.TEST_BATCH_END:
            self(context)
    def on_test_loader_end(self, context: PhaseContext) -> None:
        if self.phase == Phase.TEST_END:
            self(context)
    def on_training_end(self, context: PhaseContext) -> None:
        if self.phase == Phase.POST_TRAINING:
            self(context)
88

class PhaseContext:
    Represents the input for phase callbacks, and is constantly updated after callback calls.
    def __init__(
        self,
        epoch=None,
        batch_idx=None,
        optimizer=None,
        metrics_dict=None,
        inputs=None,
        preds=None,
        target=None,
        metrics_compute_fn=None,
        loss_avg_meter=None,
        loss_log_items=None,
        criterion=None,
        device=None,
        experiment_name=None,
        ckpt_dir=None,
        net=None,
        lr_warmup_epochs=None,
        sg_logger=None,
        train_loader=None,
        valid_loader=None,
        training_params=None,
        ddp_silent_mode=None,
        checkpoint_params=None,
        architecture=None,
        arch_params=None,
        metric_to_watch=None,
        valid_metrics=None,
        ema_model=None,
        self.epoch = epoch
        self.batch_idx = batch_idx
        self.optimizer = optimizer
        self.inputs = inputs
        self.preds = preds
        self.target = target
        self.metrics_dict = metrics_dict
        self.metrics_compute_fn = metrics_compute_fn
        self.loss_avg_meter = loss_avg_meter
        self.loss_log_items = loss_log_items
        self.criterion = criterion
        self.device = device
        self.stop_training = False
        self.experiment_name = experiment_name
        self.ckpt_dir = ckpt_dir
        self.net = net
        self.lr_warmup_epochs = lr_warmup_epochs
        self.sg_logger = sg_logger
        self.train_loader = train_loader
        self.valid_loader = valid_loader
        self.training_params = training_params
        self.ddp_silent_mode = ddp_silent_mode
        self.checkpoint_params = checkpoint_params
        self.architecture = architecture
        self.arch_params = arch_params
        self.metric_to_watch = metric_to_watch
        self.valid_metrics = valid_metrics
        self.ema_model = ema_model
    def update_context(self, **kwargs):
        for attr, attr_val in kwargs.items():
            setattr(self, attr, attr_val)
377

@register_lr_warmup(LRWarmups.LINEAR_BATCH_STEP)
class BatchStepLinearWarmupLRCallback(Callback):
    LR scheduling callback for linear step warmup on each batch step.
    LR climbs from warmup_initial_lr with to initial lr.
    def __init__(
        self,
        warmup_initial_lr: float,
        initial_lr: float,
        train_loader_len: int,
        update_param_groups: bool,
        lr_warmup_steps: int,
        training_params,
        net,
        **kwargs,
        :param warmup_initial_lr: Starting learning rate
        :param initial_lr: Target learning rate after warmup
        :param train_loader_len: Length of train data loader
        :param lr_warmup_steps: Optional. If passed, will use fixed number of warmup steps to warmup LR. Default is None.
        :param kwargs:
        super(BatchStepLinearWarmupLRCallback, self).__init__()
        if lr_warmup_steps > train_loader_len:
            logger.warning(
                f"Number of warmup steps ({lr_warmup_steps}) is greater than number of steps in epoch ({train_loader_len}). "
                f"Warmup steps will be capped to number of steps in epoch to avoid interfering with any pre-epoch LR schedulers."
        lr_warmup_steps = min(lr_warmup_steps, train_loader_len)
        learning_rates = np.linspace(start=warmup_initial_lr, stop=initial_lr, num=lr_warmup_steps, endpoint=True)
        self.lr = initial_lr
        self.initial_lr = initial_lr
        self.update_param_groups = update_param_groups
        self.training_params = training_params
        self.net = net
        self.learning_rates = learning_rates
        self.train_loader_len = train_loader_len
        self.lr_warmup_steps = lr_warmup_steps
    def on_train_batch_start(self, context: PhaseContext) -> None:
        global_training_step = context.batch_idx + context.epoch * self.train_loader_len
        if global_training_step < self.lr_warmup_steps:
            self.lr = float(self.learning_rates[global_training_step])
            self.update_lr(context.optimizer, context.epoch, context.batch_idx)
    def update_lr(self, optimizer, epoch, batch_idx=None):
        Same as in LRCallbackBase
        :param optimizer:
        :param epoch:
        :param batch_idx:
        :return:
        if self.update_param_groups:
            param_groups = unwrap_model(self.net).update_param_groups(
                optimizer.param_groups, self.lr, epoch, batch_idx, self.training_params, self.train_loader_len
            optimizer.param_groups = param_groups
        else:
            # UPDATE THE OPTIMIZERS PARAMETER
            for param_group in optimizer.param_groups:
                param_group["lr"] = self.lr
          
Optional. If passed, will use fixed number of warmup steps to warmup LR. Default is None.
              required
          kwargs
    :param warmup_initial_lr: Starting learning rate
    :param initial_lr: Target learning rate after warmup
    :param train_loader_len: Length of train data loader
    :param lr_warmup_steps: Optional. If passed, will use fixed number of warmup steps to warmup LR. Default is None.
    :param kwargs:
    super(BatchStepLinearWarmupLRCallback, self).__init__()
    if lr_warmup_steps > train_loader_len:
        logger.warning(
            f"Number of warmup steps ({lr_warmup_steps}) is greater than number of steps in epoch ({train_loader_len}). "
            f"Warmup steps will be capped to number of steps in epoch to avoid interfering with any pre-epoch LR schedulers."
    lr_warmup_steps = min(lr_warmup_steps, train_loader_len)
    learning_rates = np.linspace(start=warmup_initial_lr, stop=initial_lr, num=lr_warmup_steps, endpoint=True)
    self.lr = initial_lr
    self.initial_lr = initial_lr
    self.update_param_groups = update_param_groups
    self.training_params = training_params
    self.net = net
    self.learning_rates = learning_rates
    self.train_loader_len = train_loader_len
    self.lr_warmup_steps = lr_warmup_steps
377

def update_lr(self, optimizer, epoch, batch_idx=None):
    Same as in LRCallbackBase
    :param optimizer:
    :param epoch:
    :param batch_idx:
    :return:
    if self.update_param_groups:
        param_groups = unwrap_model(self.net).update_param_groups(
            optimizer.param_groups, self.lr, epoch, batch_idx, self.training_params, self.train_loader_len
        optimizer.param_groups = param_groups
    else:
        # UPDATE THE OPTIMIZERS PARAMETER
        for param_group in optimizer.param_groups:
            param_group["lr"] = self.lr
      
A callback that adds a visualization of a batch of segmentation predictions to context.sg_logger
  
Parameters:
        Description
        Default
667

class BinarySegmentationVisualizationCallback(PhaseCallback):
    A callback that adds a visualization of a batch of segmentation predictions to context.sg_logger
    :param phase:                   When to trigger the callback.
    :param freq:                    Frequency (in epochs) to perform this callback.
    :param batch_idx:               Batch index to perform visualization for.
    :param last_img_idx_in_batch:   Last image index to add to log. (default=-1, will take entire batch).
    def __init__(self, phase: Phase, freq: int, batch_idx: int = 0, last_img_idx_in_batch: int = -1):
        super(BinarySegmentationVisualizationCallback, self).__init__(phase)
        self.freq = freq
        self.batch_idx = batch_idx
        self.last_img_idx_in_batch = last_img_idx_in_batch
    def __call__(self, context: PhaseContext):
        if context.epoch % self.freq == 0 and context.batch_idx == self.batch_idx:
            if isinstance(context.preds, tuple):
                preds = context.preds[0].clone()
            else:
                preds = context.preds.clone()
            batch_imgs = BinarySegmentationVisualization.visualize_batch(context.inputs, preds, context.target, self.batch_idx)
            batch_imgs = [cv2.cvtColor(image, cv2.COLOR_BGR2RGB) for image in batch_imgs]
            batch_imgs = np.stack(batch_imgs)
            tag = "batch_" + str(self.batch_idx) + "_images"
            context.sg_logger.add_images(tag=tag, images=batch_imgs[: self.last_img_idx_in_batch], global_step=context.epoch, data_format="NHWC")
490

@register_lr_scheduler(LRSchedulers.COSINE)
class CosineLRCallback(LRCallbackBase):
    Hard coded step Cosine anealing learning rate scheduling.
    def __init__(self, max_epochs, cosine_final_lr_ratio, **kwargs):
        super(CosineLRCallback, self).__init__(Phase.TRAIN_BATCH_STEP, **kwargs)
        self.max_epochs = max_epochs
        self.cosine_final_lr_ratio = cosine_final_lr_ratio
    def perform_scheduling(self, context):
        effective_epoch = context.epoch - self.training_params.lr_warmup_epochs
        effective_max_epochs = self.max_epochs - self.training_params.lr_warmup_epochs - self.training_params.lr_cooldown_epochs
        current_iter = max(0, self.train_loader_len * effective_epoch + context.batch_idx - self.training_params.lr_warmup_steps)
        max_iter = self.train_loader_len * effective_max_epochs - self.training_params.lr_warmup_steps
        lr = self.compute_learning_rate(current_iter, max_iter, self.initial_lr, self.cosine_final_lr_ratio)
        self.lr = float(lr)
        self.update_lr(context.optimizer, context.epoch, context.batch_idx)
    def is_lr_scheduling_enabled(self, context):
        # Account of per-step warmup
        if self.training_params.lr_warmup_steps > 0:
            current_step = self.train_loader_len * context.epoch + context.batch_idx
            return current_step >= self.training_params.lr_warmup_steps
        post_warmup_epochs = self.training_params.max_epochs - self.training_params.lr_cooldown_epochs
        return self.training_params.lr_warmup_epochs <= context.epoch < post_warmup_epochs
    @classmethod
    def compute_learning_rate(cls, step: Union[float, np.ndarray], total_steps: float, initial_lr: float, final_lr_ratio: float):
        # the cosine starts from initial_lr and reaches initial_lr * cosine_final_lr_ratio in last epoch
        lr = 0.5 * initial_lr * (1.0 + np.cos(step / (total_steps + 1) * math.pi))
        return lr * (1 - final_lr_ratio) + (initial_lr * final_lr_ratio)
221

@register_callback(Callbacks.DECI_LAB_UPLOAD)
class DeciLabUploadCallback(PhaseCallback):
    Post-training callback for uploading and optimizing a model.
    :param model_meta_data:             Model's meta-data object. Type: ModelMetadata
    :param optimization_request_form:   Optimization request form object. Type: OptimizationRequestForm
    :param ckpt_name:                   Checkpoint filename, inside the checkpoint directory.
    def __init__(
        self,
        model_name: str,
        input_dimensions: Sequence[int],
        target_hardware_types: "Optional[List[str]]" = None,
        target_batch_size: "Optional[int]" = None,
        target_quantization_level: "Optional[str]" = None,
        ckpt_name: str = "ckpt_best.pth",
        **kwargs,
        super().__init__(phase=Phase.POST_TRAINING)
        self.input_dimensions = input_dimensions
        self.model_name = model_name
        self.target_hardware_types = target_hardware_types
        self.target_batch_size = target_batch_size
        self.target_quantization_level = target_quantization_level
        self.ckpt_name = ckpt_name
        self.platform_client = DeciClient()
    @staticmethod
    def log_optimization_failed():
        logger.info("We couldn't finish your model optimization. Visit https://console.deci.ai for details")
    def upload_model(self, model):
        This function will upload the trained model to the Deci Lab
        :param model: The resulting model from the training process
        self.platform_client.upload_model(
            model=model,
            name=self.model_name,
            input_dimensions=self.input_dimensions,
            target_hardware_types=self.target_hardware_types,
            target_batch_size=self.target_batch_size,
            target_quantization_level=self.target_quantization_level,
    def get_optimization_status(self, optimized_model_name: str):
        This function will do fetch the optimized version of the trained model and check on its benchmark status.
        The status will be checked against the server every 30 seconds and the process will timeout after 30 minutes
        or log about the successful optimization - whichever happens first.
        :param optimized_model_name: Optimized model name
        :return: Whether or not the optimized model has been benchmarked
        def handler(_signum, _frame):
            logger.error("Process timed out. Visit https://console.deci.ai for details")
            return False
        signal.signal(signal.SIGALRM, handler)
        signal.alarm(1800)
        finished = False
        while not finished:
            if self.platform_client.is_model_benchmarking(name=optimized_model_name):
                time.sleep(30)
            else:
                finished = True
        signal.alarm(0)
        return True
    def __call__(self, context: PhaseContext) -> None:
        This function will attempt to upload the trained model and schedule an optimization for it.
        :param context: Training phase context
        try:
            model = copy.deepcopy(unwrap_model(context.net))
            model_state_dict_path = os.path.join(context.ckpt_dir, self.ckpt_name)
            model_state_dict = torch.load(model_state_dict_path)["net"]
            model.load_state_dict(state_dict=model_state_dict)
            model = model.cpu()
            if hasattr(model, "prep_model_for_conversion"):
                model.prep_model_for_conversion(input_size=self.input_dimensions)
            self.upload_model(model=model)
            model_name = self.model_name
            logger.info(f"Successfully added {model_name} to the model repository")
            optimized_model_name = f"{model_name}_1_1"
            logger.info("We'll wait for the scheduled optimization to finish. Please don't close this window")
            success = self.get_optimization_status(optimized_model_name=optimized_model_name)
            if success:
                logger.info("Successfully finished your model optimization. Visit https://console.deci.ai for details")
            else:
                DeciLabUploadCallback.log_optimization_failed()
        except Exception as ex:
            DeciLabUploadCallback.log_optimization_failed()
            logger.error(ex)
      
This function will attempt to upload the trained model and schedule an optimization for it.
  
Parameters:
        Description
        Default
221

def __call__(self, context: PhaseContext) -> None:
    This function will attempt to upload the trained model and schedule an optimization for it.
    :param context: Training phase context
    try:
        model = copy.deepcopy(unwrap_model(context.net))
        model_state_dict_path = os.path.join(context.ckpt_dir, self.ckpt_name)
        model_state_dict = torch.load(model_state_dict_path)["net"]
        model.load_state_dict(state_dict=model_state_dict)
        model = model.cpu()
        if hasattr(model, "prep_model_for_conversion"):
            model.prep_model_for_conversion(input_size=self.input_dimensions)
        self.upload_model(model=model)
        model_name = self.model_name
        logger.info(f"Successfully added {model_name} to the model repository")
        optimized_model_name = f"{model_name}_1_1"
        logger.info("We'll wait for the scheduled optimization to finish. Please don't close this window")
        success = self.get_optimization_status(optimized_model_name=optimized_model_name)
        if success:
            logger.info("Successfully finished your model optimization. Visit https://console.deci.ai for details")
        else:
            DeciLabUploadCallback.log_optimization_failed()
    except Exception as ex:
        DeciLabUploadCallback.log_optimization_failed()
        logger.error(ex)
      
This function will do fetch the optimized version of the trained model and check on its benchmark status.
The status will be checked against the server every 30 seconds and the process will timeout after 30 minutes
or log about the successful optimization - whichever happens first.
  
Parameters:
        Description
        Default
190

def get_optimization_status(self, optimized_model_name: str):
    This function will do fetch the optimized version of the trained model and check on its benchmark status.
    The status will be checked against the server every 30 seconds and the process will timeout after 30 minutes
    or log about the successful optimization - whichever happens first.
    :param optimized_model_name: Optimized model name
    :return: Whether or not the optimized model has been benchmarked
    def handler(_signum, _frame):
        logger.error("Process timed out. Visit https://console.deci.ai for details")
        return False
    signal.signal(signal.SIGALRM, handler)
    signal.alarm(1800)
    finished = False
    while not finished:
        if self.platform_client.is_model_benchmarking(name=optimized_model_name):
            time.sleep(30)
        else:
            finished = True
    signal.alarm(0)
    return True
162

def upload_model(self, model):
    This function will upload the trained model to the Deci Lab
    :param model: The resulting model from the training process
    self.platform_client.upload_model(
        model=model,
        name=self.model_name,
        input_dimensions=self.input_dimensions,
        target_hardware_types=self.target_hardware_types,
        target_batch_size=self.target_batch_size,
        target_quantization_level=self.target_quantization_level,
      
A callback that adds a visualization of a batch of detection predictions to context.sg_logger
  
Parameters:
        Description
        Default
638

@register_callback(Callbacks.DETECTION_VISUALIZATION_CALLBACK)
class DetectionVisualizationCallback(PhaseCallback):
    A callback that adds a visualization of a batch of detection predictions to context.sg_logger
    :param phase:                   When to trigger the callback.
    :param freq:                    Frequency (in epochs) to perform this callback.
    :param batch_idx:               Batch index to perform visualization for.
    :param classes:                 Class list of the dataset.
    :param last_img_idx_in_batch:   Last image index to add to log. (default=-1, will take entire batch).
    def __init__(
        self,
        phase: Phase,
        freq: int,
        post_prediction_callback: DetectionPostPredictionCallback,
        classes: list,
        batch_idx: int = 0,
        last_img_idx_in_batch: int = -1,
        super(DetectionVisualizationCallback, self).__init__(phase)
        self.freq = freq
        self.post_prediction_callback = post_prediction_callback
        self.batch_idx = batch_idx
        self.classes = classes
        self.last_img_idx_in_batch = last_img_idx_in_batch
    def __call__(self, context: PhaseContext):
        if context.epoch % self.freq == 0 and context.batch_idx == self.batch_idx:
            # SOME CALCULATIONS ARE IN-PLACE IN NMS, SO CLONE THE PREDICTIONS
            preds = (context.preds[0].clone(), None)
            preds = self.post_prediction_callback(preds)
            batch_imgs = DetectionVisualization.visualize_batch(context.inputs, preds, context.target, self.batch_idx, self.classes)
            batch_imgs = [cv2.cvtColor(image, cv2.COLOR_BGR2RGB) for image in batch_imgs]
            batch_imgs = np.stack(batch_imgs)
            tag = "batch_" + str(self.batch_idx) + "_images"
            context.sg_logger.add_images(tag=tag, images=batch_imgs[: self.last_img_idx_in_batch], global_step=context.epoch, data_format="NHWC")
      
LR scheduling callback for linear step warmup. This scheduler uses a whole epoch as single step.
LR climbs from warmup_initial_lr with even steps to initial lr. When warmup_initial_lr is None - LR climb starts from
 initial_lr/(1+warmup_epochs).
          
Source code in V3_2/src/super_gradients/training/utils/callbacks/callbacks.py
293

@register_lr_warmup(LRWarmups.LINEAR_EPOCH_STEP)
class EpochStepWarmupLRCallback(LRCallbackBase):
    LR scheduling callback for linear step warmup. This scheduler uses a whole epoch as single step.
    LR climbs from warmup_initial_lr with even steps to initial lr. When warmup_initial_lr is None - LR climb starts from
     initial_lr/(1+warmup_epochs).
    def __init__(self, **kwargs):
        super(EpochStepWarmupLRCallback, self).__init__(Phase.TRAIN_EPOCH_START, **kwargs)
        self.warmup_initial_lr = self.training_params.warmup_initial_lr or self.initial_lr / (self.training_params.lr_warmup_epochs + 1)
        self.warmup_step_size = (
            (self.initial_lr - self.warmup_initial_lr) / self.training_params.lr_warmup_epochs if self.training_params.lr_warmup_epochs > 0 else 0
    def perform_scheduling(self, context):
        self.lr = self.warmup_initial_lr + context.epoch * self.warmup_step_size
        self.update_lr(context.optimizer, context.epoch, None)
    def is_lr_scheduling_enabled(self, context):
        return self.training_params.lr_warmup_epochs > 0 and self.training_params.lr_warmup_epochs >= context.epoch
429

@register_lr_scheduler(LRSchedulers.EXP)
class ExponentialLRCallback(LRCallbackBase):
    Exponential decay learning rate scheduling. Decays the learning rate by `lr_decay_factor` every epoch.
    def __init__(self, lr_decay_factor: float, **kwargs):
        super().__init__(phase=Phase.TRAIN_BATCH_STEP, **kwargs)
        self.lr_decay_factor = lr_decay_factor
    def perform_scheduling(self, context):
        effective_epoch = context.epoch - self.training_params.lr_warmup_epochs
        current_iter = self.train_loader_len * effective_epoch + context.batch_idx
        self.lr = self.initial_lr * self.lr_decay_factor ** (current_iter / self.train_loader_len)
        self.update_lr(context.optimizer, context.epoch, context.batch_idx)
    def is_lr_scheduling_enabled(self, context):
        post_warmup_epochs = self.training_params.max_epochs - self.training_params.lr_cooldown_epochs
        return self.training_params.lr_warmup_epochs <= context.epoch < post_warmup_epochs
520

@register_lr_scheduler(LRSchedulers.FUNCTION)
class FunctionLRCallback(LRCallbackBase):
    Hard coded rate scheduling for user defined lr scheduling function.
    @deprecated(version="3.2.0", reason="This callback is deprecated and will be removed in future versions.")
    def __init__(self, max_epochs, lr_schedule_function, **kwargs):
        super(FunctionLRCallback, self).__init__(Phase.TRAIN_BATCH_STEP, **kwargs)
        assert callable(lr_schedule_function), "self.lr_function must be callable"
        self.lr_schedule_function = lr_schedule_function
        self.max_epochs = max_epochs
    def is_lr_scheduling_enabled(self, context):
        post_warmup_epochs = self.training_params.max_epochs - self.training_params.lr_cooldown_epochs
        return self.training_params.lr_warmup_epochs <= context.epoch < post_warmup_epochs
    def perform_scheduling(self, context):
        effective_epoch = context.epoch - self.training_params.lr_warmup_epochs
        effective_max_epochs = self.max_epochs - self.training_params.lr_warmup_epochs - self.training_params.lr_cooldown_epochs
        self.lr = self.lr_schedule_function(
            initial_lr=self.initial_lr,
            epoch=effective_epoch,
            iter=context.batch_idx,
            max_epoch=effective_max_epochs,
            iters_per_epoch=self.train_loader_len,
        self.update_lr(context.optimizer, context.epoch, context.batch_idx)
532

class IllegalLRSchedulerMetric(Exception):
    """Exception raised illegal combination of training parameters.
    :param metric_name: Name of the metric that is not supported.
    :param metrics_dict: Dictionary of metrics that are supported.
    def __init__(self, metric_name: str, metrics_dict: dict):
        self.message = "Illegal metric name: " + metric_name + ". Expected one of metics_dics keys: " + str(metrics_dict.keys())
        super().__init__(self.message)
269

@register_callback(Callbacks.LR_CALLBACK_BASE)
class LRCallbackBase(PhaseCallback):
    Base class for hard coded learning rate scheduling regimes, implemented as callbacks.
    def __init__(self, phase, initial_lr, update_param_groups, train_loader_len, net, training_params, **kwargs):
        super(LRCallbackBase, self).__init__(phase)
        self.initial_lr = initial_lr
        self.lr = initial_lr
        self.update_param_groups = update_param_groups
        self.train_loader_len = train_loader_len
        self.net = net
        self.training_params = training_params
    def __call__(self, context: PhaseContext, **kwargs):
        if self.is_lr_scheduling_enabled(context):
            self.perform_scheduling(context)
    def is_lr_scheduling_enabled(self, context: PhaseContext):
        Predicate that controls whether to perform lr scheduling based on values in context.
        :param context: PhaseContext: current phase's context.
        :return: bool, whether to apply lr scheduling or not.
        raise NotImplementedError
    def perform_scheduling(self, context: PhaseContext):
        Performs lr scheduling based on values in context.
        :param context: PhaseContext: current phase's context.
        raise NotImplementedError
    def update_lr(self, optimizer, epoch, batch_idx=None):
        if self.update_param_groups:
            param_groups = unwrap_model(self.net).update_param_groups(
                optimizer.param_groups, self.lr, epoch, batch_idx, self.training_params, self.train_loader_len
            optimizer.param_groups = param_groups
        else:
            # UPDATE THE OPTIMIZERS PARAMETER
            for param_group in optimizer.param_groups:
                param_group["lr"] = self.lr
250

def is_lr_scheduling_enabled(self, context: PhaseContext):
    Predicate that controls whether to perform lr scheduling based on values in context.
    :param context: PhaseContext: current phase's context.
    :return: bool, whether to apply lr scheduling or not.
    raise NotImplementedError
258

def perform_scheduling(self, context: PhaseContext):
    Performs lr scheduling based on values in context.
    :param context: PhaseContext: current phase's context.
    raise NotImplementedError
      
Learning rate scheduler callback.
When passing call a metrics_dict, with a key=self.metric_name, the value of that metric will monitored
     for ReduceLROnPlateau (i.e step(metrics_dict[self.metric_name]).
  
Parameters:
        Description
        Default
                torch.optim.lr_scheduler._LRScheduler
          
Learning rate scheduler to be called step() with.
              required
          metric_name
          
Metric name for ReduceLROnPlateau learning rate scheduler.
          phase
                Phase
          
Phase of when to trigger it.
              required
563

@register_callback(Callbacks.LR_SCHEDULER)
class LRSchedulerCallback(PhaseCallback):
    Learning rate scheduler callback.
    When passing __call__ a metrics_dict, with a key=self.metric_name, the value of that metric will monitored
         for ReduceLROnPlateau (i.e step(metrics_dict[self.metric_name]).
    :param scheduler:       Learning rate scheduler to be called step() with.
    :param metric_name:     Metric name for ReduceLROnPlateau learning rate scheduler.
    :param phase:           Phase of when to trigger it.
    def __init__(self, scheduler: torch.optim.lr_scheduler._LRScheduler, phase: Phase, metric_name: str = None):
        super(LRSchedulerCallback, self).__init__(phase)
        self.scheduler = scheduler
        self.metric_name = metric_name
    def __call__(self, context: PhaseContext):
        if context.lr_warmup_epochs <= context.epoch:
            if self.metric_name and self.metric_name in context.metrics_dict.keys():
                self.scheduler.step(context.metrics_dict[self.metric_name])
            elif self.metric_name is None:
                self.scheduler.step()
            else:
                raise IllegalLRSchedulerMetric(self.metric_name, context.metrics_dict)
    def __repr__(self):
        return "LRSchedulerCallback: " + repr(self.scheduler)
305

@register_lr_warmup(LRWarmups.LINEAR_STEP)
class LinearStepWarmupLRCallback(EpochStepWarmupLRCallback):
    """Deprecated, use EpochStepWarmupLRCallback instead"""
    def __init__(self, **kwargs):
        logger.warning(
            f"Parameter {LRWarmups.LINEAR_STEP} has been made deprecated and will be removed in the next SG release. "
            f"Please use `{LRWarmups.LINEAR_EPOCH_STEP}` instead."
        super(LinearStepWarmupLRCallback, self).__init__(**kwargs)
      
Pre-training callback that verifies model conversion to onnx given specified conversion parameters.
The model is converted, then inference is applied with onnx runtime.
Use this callback with the same args as DeciPlatformCallback to prevent conversion fails at the end of training.
  
Parameters:
        Description
        Default
113

@register_callback(Callbacks.MODEL_CONVERSION_CHECK)
class ModelConversionCheckCallback(PhaseCallback):
    Pre-training callback that verifies model conversion to onnx given specified conversion parameters.
    The model is converted, then inference is applied with onnx runtime.
    Use this callback with the same args as DeciPlatformCallback to prevent conversion fails at the end of training.
    :param model_name:              Model's name
    :param input_dimensions:        Model's input dimensions
    :param primary_batch_size:      Model's primary batch size
    :param opset_version:           (default=11)
    :param do_constant_folding:     (default=True)
    :param dynamic_axes:            (default={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})
    :param input_names:             (default=["input"])
    :param output_names:            (default=["output"])
    :param rtol:                    (default=1e-03)
    :param atol:                    (default=1e-05)
    def __init__(self, model_name: str, input_dimensions: Sequence[int], primary_batch_size: int, **kwargs):
        super(ModelConversionCheckCallback, self).__init__(phase=Phase.PRE_TRAINING)
        self.model_name = model_name
        self.input_dimensions = input_dimensions
        self.primary_batch_size = primary_batch_size
        self.opset_version = kwargs.get("opset_version", 10)
        self.do_constant_folding = kwargs.get("do_constant_folding", None) if kwargs.get("do_constant_folding", None) else True
        self.input_names = kwargs.get("input_names") or ["input"]
        self.output_names = kwargs.get("output_names") or ["output"]
        self.dynamic_axes = kwargs.get("dynamic_axes") or {"input": {0: "batch_size"}, "output": {0: "batch_size"}}
        self.rtol = kwargs.get("rtol", 1e-03)
        self.atol = kwargs.get("atol", 1e-05)
    def __call__(self, context: PhaseContext):
        model = copy.deepcopy(unwrap_model(context.net))
        model = model.cpu()
        model.eval()  # Put model into eval mode
        if hasattr(model, "prep_model_for_conversion"):
            model.prep_model_for_conversion(input_size=self.input_dimensions)
        x = torch.randn(self.primary_batch_size, *self.input_dimensions, requires_grad=False)
        tmp_model_path = os.path.join(context.ckpt_dir, self.model_name + "_tmp.onnx")
        with torch.no_grad():
            torch_out = model(x)
        torch.onnx.export(
            model,  # Model being run
            x,  # Model input (or a tuple for multiple inputs)
            tmp_model_path,  # Where to save the model (can be a file or file-like object)
            export_params=True,  # Store the trained parameter weights inside the model file
            opset_version=self.opset_version,
            do_constant_folding=self.do_constant_folding,
            input_names=self.input_names,
            output_names=self.output_names,
            dynamic_axes=self.dynamic_axes,
        onnx_model = onnx.load(tmp_model_path)
        onnx.checker.check_model(onnx_model)
        ort_session = onnxruntime.InferenceSession(tmp_model_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
        # compute ONNX Runtime output prediction
        ort_inputs = {ort_session.get_inputs()[0].name: x.cpu().numpy()}
        ort_outs = ort_session.run(None, ort_inputs)
        # TODO: Ideally we don't want to check this but have the certainty of just calling torch_out.cpu()
        if isinstance(torch_out, List) or isinstance(torch_out, tuple):
            torch_out = torch_out[0]
        # compare ONNX Runtime and PyTorch results
        np.testing.assert_allclose(torch_out.cpu().numpy(), ort_outs[0], rtol=self.rtol, atol=self.atol)
        os.remove(tmp_model_path)
        logger.info("Exported model has been tested with ONNXRuntime, and the result looks good!")
598

class PhaseContextTestCallback(PhaseCallback):
    A callback that saves the phase context the for testing.
    def __init__(self, phase: Phase):
        super(PhaseContextTestCallback, self).__init__(phase)
        self.context = None
    def __call__(self, context: PhaseContext):
        self.context = context
452

@register_lr_scheduler(LRSchedulers.POLY)
class PolyLRCallback(LRCallbackBase):
    Hard coded polynomial decay learning rate scheduling (i.e at specific milestones).
    def __init__(self, max_epochs, **kwargs):
        super(PolyLRCallback, self).__init__(Phase.TRAIN_BATCH_STEP, **kwargs)
        self.max_epochs = max_epochs
    def perform_scheduling(self, context):
        effective_epoch = context.epoch - self.training_params.lr_warmup_epochs
        effective_max_epochs = self.max_epochs - self.training_params.lr_warmup_epochs - self.training_params.lr_cooldown_epochs
        current_iter = (self.train_loader_len * effective_epoch + context.batch_idx) / self.training_params.batch_accumulate
        max_iter = self.train_loader_len * effective_max_epochs / self.training_params.batch_accumulate
        self.lr = self.initial_lr * pow((1.0 - (current_iter / max_iter)), 0.9)
        self.update_lr(context.optimizer, context.epoch, context.batch_idx)
    def is_lr_scheduling_enabled(self, context):
        post_warmup_epochs = self.training_params.max_epochs - self.training_params.lr_cooldown_epochs
        return self.training_params.lr_warmup_epochs <= context.epoch < post_warmup_epochs
742

@register_callback(Callbacks.ROBOFLOW_RESULT_CALLBACK)
class RoboflowResultCallback(Callback):
    """Append the training results to a csv file. Be aware that this does not fully overwrite the existing file, just appends."""
    def __init__(self, dataset_name: str, output_path: Optional[str] = None):
        :param dataset_name:    Name of the dataset that was used to train the model.
        :param output_path:     Full path to the output csv file. By default, save at 'checkpoint_dir/results.csv'
        self.dataset_name = dataset_name
        self.output_path = output_path or os.path.join(get_project_checkpoints_dir_path(), "results.csv")
        if self.output_path is None:
            raise ValueError("Output path must be specified")
        super(RoboflowResultCallback, self).__init__()
    @multi_process_safe
    def on_training_end(self, context: PhaseContext):
        with open(self.output_path, mode="a", newline="") as csv_file:
            writer = csv.writer(csv_file)
            mAP = context.metrics_dict["[email protected]:0.95"].item()
            writer.writerow([self.dataset_name, mAP])
                Optional[str]
          
Full path to the output csv file. By default, save at 'checkpoint_dir/results.csv'
734

def __init__(self, dataset_name: str, output_path: Optional[str] = None):
    :param dataset_name:    Name of the dataset that was used to train the model.
    :param output_path:     Full path to the output csv file. By default, save at 'checkpoint_dir/results.csv'
    self.dataset_name = dataset_name
    self.output_path = output_path or os.path.join(get_project_checkpoints_dir_path(), "results.csv")
    if self.output_path is None:
        raise ValueError("Output path must be specified")
    super(RoboflowResultCallback, self).__init__()
408

@register_lr_scheduler(LRSchedulers.STEP)
class StepLRCallback(LRCallbackBase):
    Hard coded step learning rate scheduling (i.e at specific milestones).
    def __init__(self, lr_updates, lr_decay_factor, step_lr_update_freq=None, **kwargs):
        super(StepLRCallback, self).__init__(Phase.TRAIN_EPOCH_END, **kwargs)
        if step_lr_update_freq and len(lr_updates):
            raise ValueError("Only one of [lr_updates, step_lr_update_freq] should be passed to StepLRCallback constructor")
        if step_lr_update_freq:
            max_epochs = self.training_params.max_epochs - self.training_params.lr_cooldown_epochs
            warmup_epochs = self.training_params.lr_warmup_epochs
            lr_updates = [
                int(np.ceil(step_lr_update_freq * x)) for x in range(1, max_epochs) if warmup_epochs <= int(np.ceil(step_lr_update_freq * x)) < max_epochs
        elif self.training_params.lr_cooldown_epochs > 0:
            logger.warning("Specific lr_updates were passed along with cooldown_epochs > 0," " cooldown will have no effect.")
        self.lr_updates = lr_updates
        self.lr_decay_factor = lr_decay_factor
    def perform_scheduling(self, context):
        num_updates_passed = [x for x in self.lr_updates if x <= context.epoch]
        self.lr = self.initial_lr * self.lr_decay_factor ** len(num_updates_passed)
        self.update_lr(context.optimizer, context.epoch, None)
    def is_lr_scheduling_enabled(self, context):
        return self.training_params.lr_warmup_epochs <= context.epoch
      
Phase callback that collects the learning rates in lr_placeholder at the end of each epoch (used for testing). In
 the case of multiple parameter groups (i.e multiple learning rates) the learning rate is collected from the first
 one. The phase is VALIDATION_EPOCH_END to ensure all lr updates have been performed before calling this callback.
          
Source code in V3_2/src/super_gradients/training/utils/callbacks/callbacks.py
757

class TestLRCallback(PhaseCallback):
    Phase callback that collects the learning rates in lr_placeholder at the end of each epoch (used for testing). In
     the case of multiple parameter groups (i.e multiple learning rates) the learning rate is collected from the first
     one. The phase is VALIDATION_EPOCH_END to ensure all lr updates have been performed before calling this callback.
    def __init__(self, lr_placeholder):
        super(TestLRCallback, self).__init__(Phase.VALIDATION_EPOCH_END)
        self.lr_placeholder = lr_placeholder
    def __call__(self, context: PhaseContext):
        self.lr_placeholder.append(context.optimizer.param_groups[0]["lr"])
      
TrainingStageSwitchCallback
A phase callback that is called at a specific epoch (epoch start) to support multi-stage training.
It does so by manipulating the objects inside the context.
  
Parameters:
        Description
        Default
695

class TrainingStageSwitchCallbackBase(PhaseCallback):
    TrainingStageSwitchCallback
    A phase callback that is called at a specific epoch (epoch start) to support multi-stage training.
    It does so by manipulating the objects inside the context.
    :param next_stage_start_epoch: Epoch idx to apply the stage change.
    def __init__(self, next_stage_start_epoch: int):
        super(TrainingStageSwitchCallbackBase, self).__init__(phase=Phase.TRAIN_EPOCH_START)
        self.next_stage_start_epoch = next_stage_start_epoch
    def __call__(self, context: PhaseContext):
        if context.epoch == self.next_stage_start_epoch:
            self.apply_stage_change(context)
    def apply_stage_change(self, context: PhaseContext):
        This method is called when the callback is fired on the next_stage_start_epoch,
         and holds the stage change logic that should be applied to the context's objects.
        :param context: PhaseContext, context of current phase
        raise NotImplementedError
      
This method is called when the callback is fired on the next_stage_start_epoch,
 and holds the stage change logic that should be applied to the context's objects.
  
Parameters:
        Description
        Default
695