The right place for an "essential" callback
❓ Questions and Help
What is your question?
I am currently using an ordinal loss formulation that cuts up a real-valued output space into regions using cutpoints. Each region is linked to a discrete ordered label. After the backward (optimizer step), the model requires cutpoints to be re-arranged in an ascending order. I've provided a brief snippet from the original author on how to achieve this using vanilla pytorch (library :  spacecutter)
def ascension_callback(margin=0.0, min_val=-1.0e6):

    def _clip(module):
        if isinstance(module, LogisticCumulativeLink):
            cutpoints = module.cutpoints.data
            for i in range(cutpoints.shape[0] - 1):
                cutpoints[i].clamp_(
                    min_val, cutpoints[i + 1] - margin
                )
    
    return _clip

callback = ascension_callback()

# In your training loop, do the following:

for data in data_iterator:
    # Calculate loss
    # Step optimizer
    model.apply(callback)
What have you tried?
To achieve this in pytorch-lightning, I converted the original callback code to the form below :
class AscensionCallback(Callback):
    """
    Ensure that each cutpoint is ordered in ascending value.
    e.g.

    .. < cutpoint[i - 1] < cutpoint[i] < cutpoint[i + 1] < ...

    This is done by clipping the cutpoint values at the end of a batch gradient
    update. By no means is this an efficient way to do things, but it works out
    of the box with stochastic gradient descent.

    Parameters
    ----------
    margin : float, (default=0.0)
        The minimum value between any two adjacent cutpoints.
        e.g. enforce that cutpoint[i - 1] + margin < cutpoint[i]
    min_val : float, (default=-1e6)
        Minimum value that the smallest cutpoint may take.
    """

    def __init__(self, margin: float = 0.0, min_val: float = -1.0e6) -> None:
        super().__init__()
        self.margin = margin
        self.min_val = min_val

    def clip(self, module: Module) -> None:
        # NOTE: Only works for LogisticCumulativeLink right now
        # We assume the cutpoints parameters are called `cutpoints`.
        if isinstance(module, LogisticCumulativeLink):
            cutpoints = module.cutpoints.data
            for i in range(cutpoints.shape[0] - 1):
                cutpoints[i].clamp_(self.min_val, cutpoints[i + 1] - self.margin)

    def on_batch_end(self, trainer, pl_module):
        pl_module.model.apply(self.clip)
I then call it using the trainer like so:
trainer = Trainer.from_argparse_args(args, callbacks=[AscensionCallback()])
Questions?


The documentation suggests using callbacks for non-essential code. In this case, the callback is essential and directly modifies the cutpoints. Are callbacks the right way to do something like this?


If not, how can I incorporate this into the main LightningModule such that the clip function is called after optimizer step during training. In which function do I place a call to this "AscensionCallback()"


What's your environment?

OS: Ubuntu 20.04.1 LTS
Packaging : pip
Version : 0.8.5