Skip to content

callbacks

shrinkai.distillation.callbacks

Ready-to-use training callbacks for DistillationEngine.fit / Distiller.fit.

A callback is any callable accepting (epoch: int, metrics: dict[str, float]), as already supported by fit(..., callbacks=[...]). EarlyStopping additionally exposes a stop boolean attribute: after invoking all callbacks, the training loop checks getattr(callback, "stop", False) on each of them and breaks out of the epoch loop if any callback requests it. Plain function callbacks are unaffected by this check.

Classes:

Name Description
EarlyStopping

Stops training when a monitored metric has stopped improving.
ModelCheckpoint

Saves a model's weights to disk during training.

Classes

EarlyStopping

Stops training when a monitored metric has stopped improving.

Attributes:

Name Type Description
stop bool

Set to True once patience is exhausted. Read by DistillationEngine.fit after each epoch to interrupt the loop early.

Methods:

Name Description
__call__

Updates internal state and sets self.stop if patience is exhausted.
__init__

Initializes the EarlyStopping callback.
Source code in src/shrinkai/distillation/callbacks.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
class EarlyStopping:
    """Stops training when a monitored metric has stopped improving.

    Attributes:
        stop (bool): Set to True once `patience` is exhausted. Read by
            `DistillationEngine.fit` after each epoch to interrupt the loop early.
    """

    def __init__(
        self,
        monitor: str = "val_loss",
        patience: int = 5,
        mode: Literal["min", "max"] = "min",
        min_delta: float = 0.0,
    ) -> None:
        """Initializes the EarlyStopping callback.

        Args:
            monitor: Metric key to watch in the epoch summary dict (e.g. "val_loss").
            patience: Number of consecutive non-improving epochs tolerated before
                training is stopped.
            mode: "min" if lower values of `monitor` are better, "max" otherwise.
            min_delta: Minimum absolute change to qualify as an improvement.

        Raises:
            ValueError: If `mode` is not "min" or "max".
        """
        if mode not in ("min", "max"):
            raise ValueError(f"mode must be 'min' or 'max', got '{mode}'.")

        self.monitor = monitor
        self.patience = patience
        self.mode = mode
        self.min_delta = min_delta

        self.best_score: float | None = None
        self.num_bad_epochs = 0
        self.stop = False

    def _is_improvement(self, current: float) -> bool:
        if self.best_score is None:
            return True
        if self.mode == "min":
            return current < self.best_score - self.min_delta
        return current > self.best_score + self.min_delta

    def __call__(self, epoch: int, metrics: dict[str, float]) -> None:
        """Updates internal state and sets `self.stop` if patience is exhausted.

        Args:
            epoch: Current 1-based epoch index.
            metrics: Epoch summary dict, as passed by `DistillationEngine.fit`.
        """
        if self.monitor not in metrics:
            logger.warning(
                "EarlyStopping: metric '%s' not found in epoch %d summary. Skipping check.",
                self.monitor,
                epoch,
            )
            return

        current = metrics[self.monitor]
        if self._is_improvement(current):
            self.best_score = current
            self.num_bad_epochs = 0
        else:
            self.num_bad_epochs += 1

        if self.num_bad_epochs >= self.patience:
            logger.info(
                "EarlyStopping: '%s' did not improve for %d epoch(s). Stopping at epoch %d.",
                self.monitor,
                self.patience,
                epoch,
            )
            self.stop = True
Methods:
__call__ 
__call__(epoch: int, metrics: dict[str, float]) -> None

Updates internal state and sets self.stop if patience is exhausted.

Parameters:

Name Type Description Default
epoch int

Current 1-based epoch index.

 required
metrics dict[str, float]

Epoch summary dict, as passed by DistillationEngine.fit.

 required
Source code in src/shrinkai/distillation/callbacks.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def __call__(self, epoch: int, metrics: dict[str, float]) -> None:
    """Updates internal state and sets `self.stop` if patience is exhausted.

    Args:
        epoch: Current 1-based epoch index.
        metrics: Epoch summary dict, as passed by `DistillationEngine.fit`.
    """
    if self.monitor not in metrics:
        logger.warning(
            "EarlyStopping: metric '%s' not found in epoch %d summary. Skipping check.",
            self.monitor,
            epoch,
        )
        return

    current = metrics[self.monitor]
    if self._is_improvement(current):
        self.best_score = current
        self.num_bad_epochs = 0
    else:
        self.num_bad_epochs += 1

    if self.num_bad_epochs >= self.patience:
        logger.info(
            "EarlyStopping: '%s' did not improve for %d epoch(s). Stopping at epoch %d.",
            self.monitor,
            self.patience,
            epoch,
        )
        self.stop = True
__init__ 
__init__(
    monitor: str = "val_loss",
    patience: int = 5,
    mode: Literal["min", "max"] = "min",
    min_delta: float = 0.0,
) -> None

Initializes the EarlyStopping callback.

Parameters:

Name Type Description Default
monitor str

Metric key to watch in the epoch summary dict (e.g. "val_loss").

 'val_loss'
patience int

Number of consecutive non-improving epochs tolerated before training is stopped.

 5
mode Literal['min', 'max']

"min" if lower values of monitor are better, "max" otherwise.

 'min'
min_delta float

Minimum absolute change to qualify as an improvement.

 0.0

Raises:

Type Description
ValueError

If mode is not "min" or "max".
Source code in src/shrinkai/distillation/callbacks.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def __init__(
    self,
    monitor: str = "val_loss",
    patience: int = 5,
    mode: Literal["min", "max"] = "min",
    min_delta: float = 0.0,
) -> None:
    """Initializes the EarlyStopping callback.

    Args:
        monitor: Metric key to watch in the epoch summary dict (e.g. "val_loss").
        patience: Number of consecutive non-improving epochs tolerated before
            training is stopped.
        mode: "min" if lower values of `monitor` are better, "max" otherwise.
        min_delta: Minimum absolute change to qualify as an improvement.

    Raises:
        ValueError: If `mode` is not "min" or "max".
    """
    if mode not in ("min", "max"):
        raise ValueError(f"mode must be 'min' or 'max', got '{mode}'.")

    self.monitor = monitor
    self.patience = patience
    self.mode = mode
    self.min_delta = min_delta

    self.best_score: float | None = None
    self.num_bad_epochs = 0
    self.stop = False

ModelCheckpoint

Saves a model's weights to disk during training.

Holds a direct reference to the module to save (typically the student), so it plugs into fit(callbacks=[...]) without changing the existing (epoch, metrics) -> None callback signature.

Methods:

Name Description
__call__

Saves the model's weights, respecting save_best_only.
__init__

Initializes the ModelCheckpoint callback.
Source code in src/shrinkai/distillation/callbacks.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
class ModelCheckpoint:
    """Saves a model's weights to disk during training.

    Holds a direct reference to the module to save (typically the student), so it
    plugs into `fit(callbacks=[...])` without changing the existing
    `(epoch, metrics) -> None` callback signature.
    """

    def __init__(
        self,
        model: nn.Module,
        filepath: str | Path,
        monitor: str = "val_loss",
        mode: Literal["min", "max"] = "min",
        save_best_only: bool = True,
    ) -> None:
        """Initializes the ModelCheckpoint callback.

        Args:
            model: The module whose `state_dict()` is saved (e.g. `distiller.student`).
            filepath: Destination path for the saved weights.
            monitor: Metric key to watch when `save_best_only` is True.
            mode: "min" if lower values of `monitor` are better, "max" otherwise.
            save_best_only: If True, only overwrite `filepath` when `monitor` improves
                over its best value so far. If False, save unconditionally every epoch.

        Raises:
            ValueError: If `mode` is not "min" or "max".
        """
        if mode not in ("min", "max"):
            raise ValueError(f"mode must be 'min' or 'max', got '{mode}'.")

        self.model = model
        self.filepath = Path(filepath)
        self.monitor = monitor
        self.mode = mode
        self.save_best_only = save_best_only
        self.best_score: float | None = None

    def _is_improvement(self, current: float) -> bool:
        if self.best_score is None:
            return True
        if self.mode == "min":
            return current < self.best_score
        return current > self.best_score

    def _save(self) -> None:
        self.filepath.parent.mkdir(parents=True, exist_ok=True)
        torch.save(self.model.state_dict(), self.filepath)

    def __call__(self, epoch: int, metrics: dict[str, float]) -> None:
        """Saves the model's weights, respecting `save_best_only`.

        Args:
            epoch: Current 1-based epoch index.
            metrics: Epoch summary dict, as passed by `DistillationEngine.fit`.
        """
        if not self.save_best_only:
            self._save()
            return

        if self.monitor not in metrics:
            logger.warning(
                "ModelCheckpoint: metric '%s' not found in epoch %d summary. Skipping save.",
                self.monitor,
                epoch,
            )
            return

        current = metrics[self.monitor]
        if self._is_improvement(current):
            self.best_score = current
            self._save()
            logger.info(
                "ModelCheckpoint: '%s' improved to %.4f at epoch %d. Saved to %s.",
                self.monitor,
                current,
                epoch,
                self.filepath,
            )
Methods:
__call__ 
__call__(epoch: int, metrics: dict[str, float]) -> None

Saves the model's weights, respecting save_best_only.

Parameters:

Name Type Description Default
epoch int

Current 1-based epoch index.

 required
metrics dict[str, float]

Epoch summary dict, as passed by DistillationEngine.fit.

 required
Source code in src/shrinkai/distillation/callbacks.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
def __call__(self, epoch: int, metrics: dict[str, float]) -> None:
    """Saves the model's weights, respecting `save_best_only`.

    Args:
        epoch: Current 1-based epoch index.
        metrics: Epoch summary dict, as passed by `DistillationEngine.fit`.
    """
    if not self.save_best_only:
        self._save()
        return

    if self.monitor not in metrics:
        logger.warning(
            "ModelCheckpoint: metric '%s' not found in epoch %d summary. Skipping save.",
            self.monitor,
            epoch,
        )
        return

    current = metrics[self.monitor]
    if self._is_improvement(current):
        self.best_score = current
        self._save()
        logger.info(
            "ModelCheckpoint: '%s' improved to %.4f at epoch %d. Saved to %s.",
            self.monitor,
            current,
            epoch,
            self.filepath,
        )
__init__ 
__init__(
    model: Module,
    filepath: str | Path,
    monitor: str = "val_loss",
    mode: Literal["min", "max"] = "min",
    save_best_only: bool = True,
) -> None

Initializes the ModelCheckpoint callback.

Parameters:

Name Type Description Default
model Module

The module whose state_dict() is saved (e.g. distiller.student).

 required
filepath str | Path

Destination path for the saved weights.

 required
monitor str

Metric key to watch when save_best_only is True.

 'val_loss'
mode Literal['min', 'max']

"min" if lower values of monitor are better, "max" otherwise.

 'min'
save_best_only bool

If True, only overwrite filepath when monitor improves over its best value so far. If False, save unconditionally every epoch.

 True

Raises:

Type Description
ValueError

If mode is not "min" or "max".
Source code in src/shrinkai/distillation/callbacks.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def __init__(
    self,
    model: nn.Module,
    filepath: str | Path,
    monitor: str = "val_loss",
    mode: Literal["min", "max"] = "min",
    save_best_only: bool = True,
) -> None:
    """Initializes the ModelCheckpoint callback.

    Args:
        model: The module whose `state_dict()` is saved (e.g. `distiller.student`).
        filepath: Destination path for the saved weights.
        monitor: Metric key to watch when `save_best_only` is True.
        mode: "min" if lower values of `monitor` are better, "max" otherwise.
        save_best_only: If True, only overwrite `filepath` when `monitor` improves
            over its best value so far. If False, save unconditionally every epoch.

    Raises:
        ValueError: If `mode` is not "min" or "max".
    """
    if mode not in ("min", "max"):
        raise ValueError(f"mode must be 'min' or 'max', got '{mode}'.")

    self.model = model
    self.filepath = Path(filepath)
    self.monitor = monitor
    self.mode = mode
    self.save_best_only = save_best_only
    self.best_score: float | None = None