Skip to content

Callbacks

continuiti.trainer.callbacks

Callbacks for Trainer in continuiti.

Callback

Callback base class for fit method of Trainer.

__call__(logs)

Callback function. Called at the end of each epoch.

PARAMETER DESCRIPTION
logs

Training logs.

TYPE: Logs

Source code in src/continuiti/trainer/callbacks.py 
25
26
27
28
29
30
31
def __call__(self, logs: Logs):
    """Callback function.
    Called at the end of each epoch.

    Args:
        logs: Training logs.
    """

step(logs)

Called after every gradient step.

PARAMETER DESCRIPTION
logs

Training logs.

TYPE: Logs

Source code in src/continuiti/trainer/callbacks.py 
33
34
35
36
37
38
def step(self, logs: Logs):
    """Called after every gradient step.

    Args:
        logs: Training logs.
    """

on_train_begin()

Called at the beginning of training.

Source code in src/continuiti/trainer/callbacks.py 
40
41
def on_train_begin(self):
    """Called at the beginning of training."""

on_train_end()

Called at the end of training.

Source code in src/continuiti/trainer/callbacks.py 
43
44
def on_train_end(self):
    """Called at the end of training."""

PrintTrainingLoss(epochs=None, steps=None)

Bases: Callback

Callback to print training/test loss.

PARAMETER DESCRIPTION
epochs

Number of epochs. Default is None.

TYPE: Optional[int] DEFAULT: None

steps

Number of steps per epoch. Default is None.

TYPE: Optional[int] DEFAULT: None

Source code in src/continuiti/trainer/callbacks.py 
56
57
58
59
60
61
def __init__(self, epochs: Optional[int] = None, steps: Optional[int] = None):
    self.epochs = epochs
    self.steps = steps
    self.steps_performed = 0
    self.start_time = time()
    super().__init__()

__call__(logs)

Callback function. Called at the end of each epoch.

PARAMETER DESCRIPTION
logs

Training logs.

TYPE: Logs

Source code in src/continuiti/trainer/callbacks.py 
 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
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def __call__(self, logs: Logs):
    """Callback function.
    Called at the end of each epoch.

    Args:
        logs: Training logs.
    """
    elapsed = time() - self.start_time
    sec_per_step = elapsed / self.steps_performed

    s = ""
    s += f"Epoch {logs.epoch:2.0f}"
    if self.epochs is not None:
        s += f"/{self.epochs}"
    s += "  "

    s += f"Step {logs.step:2.0f}"
    if self.steps is not None:
        s += f"/{self.steps}"

        # Progress bar
        n = 20
        done = math.ceil(logs.step / self.steps * n)
        s += "  [" + "=" * done + " " * (n - done) + "]"
    s += "  "

    if sec_per_step < 1:
        s += f"{sec_per_step * 1000:.0f}ms/step"
    else:
        s += f"{sec_per_step:.2f}s/step"

    def to_min(t):
        min = t // 60
        sec = math.floor(t) % 60
        return f"{min:.0f}:{sec:02.0f}min"

    s += f"  [{to_min(elapsed)}"

    if self.epochs is not None and self.steps is not None:
        remaining_steps = (self.epochs - logs.epoch) * self.steps
        remaining_steps += self.steps - logs.step
        eta = remaining_steps * sec_per_step
        s += f"<{to_min(eta)}"
    s += "] - "

    s += f"loss/train = {logs.loss_train:.4e}  "

    if logs.loss_test is not None:
        s += f"loss/test = {logs.loss_test:.4e} "

    print("\r" + s, end="")

step(logs)

Called after every gradient step.

PARAMETER DESCRIPTION
logs

Training logs.

TYPE: Logs

Source code in src/continuiti/trainer/callbacks.py 
115
116
117
118
119
120
121
122
def step(self, logs: Logs):
    """Called after every gradient step.

    Args:
        logs: Training logs.
    """
    self.steps_performed += 1
    self.__call__(logs)

on_train_begin()

Called at the beginning of training.

Source code in src/continuiti/trainer/callbacks.py 
124
125
126
def on_train_begin(self):
    """Called at the beginning of training."""
    self.start_time = time()

on_train_end()

Called at the end of training.

Source code in src/continuiti/trainer/callbacks.py 
128
129
130
def on_train_end(self):
    """Called at the end of training."""
    print("")

LearningCurve(keys=None)

Bases: Callback

Callback to plot learning curve.

PARAMETER DESCRIPTION
keys

List of keys to plot. Default is ["loss_train"].

TYPE: Optional[List[str]] DEFAULT: None

Source code in src/continuiti/trainer/callbacks.py 
141
142
143
144
145
146
147
def __init__(self, keys: Optional[List[str]] = None):
    if keys is None:
        keys = ["loss_train", "loss_test"]

    self.keys = keys
    self.on_train_begin()
    super().__init__()

__call__(logs)

Callback function. Called at the end of each epoch.

PARAMETER DESCRIPTION
logs

Training logs.

TYPE: Logs

Source code in src/continuiti/trainer/callbacks.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
def __call__(self, logs: Logs):
    """Callback function.
    Called at the end of each epoch.

    Args:
        logs: Training logs.
    """
    for key in self.keys:
        try:
            test = logs.__getattribute__(key)
            self.losses[key].append(test)
        except AttributeError:
            pass

on_train_begin()

Called at the beginning of training.

Source code in src/continuiti/trainer/callbacks.py 
163
164
165
def on_train_begin(self):
    """Called at the beginning of training."""
    self.losses = {key: [] for key in self.keys}

on_train_end()

Called at the end of training.

Source code in src/continuiti/trainer/callbacks.py 
167
168
169
170
171
172
173
174
175
176
177
178
def on_train_end(self):
    """Called at the end of training."""
    for key in self.keys:
        vals = self.losses[key]
        epochs = list(range(1, len(vals) + 1))
        plt.plot(epochs, vals)

    plt.yscale("log")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.legend(self.keys)
    plt.show()

OptunaCallback(trial)

Bases: Callback

Callback to report intermediate values to Optuna.

PARAMETER DESCRIPTION
trial

Optuna trial.
Source code in src/continuiti/trainer/callbacks.py 
189
190
191
def __init__(self, trial):
    self.trial = trial
    super().__init__()

__call__(logs)

Callback function. Called at the end of each epoch.

PARAMETER DESCRIPTION
logs

Training logs.

TYPE: Logs

Source code in src/continuiti/trainer/callbacks.py 
193
194
195
196
197
198
199
200
def __call__(self, logs: Logs):
    """Callback function.
    Called at the end of each epoch.

    Args:
        logs: Training logs.
    """
    self.trial.report(logs.loss_train, step=logs.epoch)

MLFlowLogger(operator=None)

Bases: Callback

Callback to log metrics to MLFlow.

PARAMETER DESCRIPTION
operator

Operator to save snapshots. Default is None.

TYPE: Optional[Operator] DEFAULT: None

Source code in src/continuiti/trainer/callbacks.py 
211
212
213
214
def __init__(self, operator: Optional[Operator] = None):
    self.operator = operator
    self.best_loss = math.inf
    super().__init__()

__call__(logs)

Callback function. Called at the end of each epoch.

PARAMETER DESCRIPTION
logs

Training logs.

TYPE: Logs

Source code in src/continuiti/trainer/callbacks.py 
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
def __call__(self, logs: Logs):
    """Callback function.
    Called at the end of each epoch.

    Args:
        logs: Training logs.
    """
    mlflow.log_metric("loss/train", logs.loss_train, step=logs.epoch)
    loss = logs.loss_train
    if logs.loss_test is not None:
        mlflow.log_metric("loss/test", logs.loss_test, step=logs.epoch)
        loss = logs.loss_test

    # Save best model
    self.best_loss = min(self.best_loss, loss)
    if self.best_loss == loss:
        self._save_model("best")

on_train_begin()

Called at the beginning of training.

Source code in src/continuiti/trainer/callbacks.py 
234
235
236
237
238
def on_train_begin(self):
    """Called at the beginning of training."""
    if not mlflow.active_run():
        mlflow.start_run()
    self._save_model("initial")

on_train_end()

Called at the end of training.

Source code in src/continuiti/trainer/callbacks.py 
240
241
242
243
def on_train_end(self):
    """Called at the end of training."""
    self._save_model("final")
    mlflow.end_run()
Last update: 2024-08-20
Created: 2024-08-20