Skip to content

myco.callbacks

Functions to report on active model training status.

CyclicLRPolicy

Bases: tf.keras.callbacks.Callback

Cyclically modify learning rates

Source code in myco/callbacks.py 
 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
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
class CyclicLRPolicy(tf.keras.callbacks.Callback):
    """Cyclically modify learning rates"""

    def __init__(
        self,
        min_lr: float = 0.0001,
        max_lr: float = None,
        lr_range_scaler: float = 10,
        mode: str = "one_cycle",
        momentum: float = 0.8,
        sample_size: int = None,
        batch_size: int = None,
        step_factor: int = 1,
    ):
        """Create a cyclic learning rate callback to modify LR by batch.

        Reference: L.N. Smith 2017 https://arxiv.org/pdf/1506.01186.pdf

        Args:
            min_lr: minimum bounding learning rate.
            max_lr: maximum bounding learning rate.
            lr_range_scaler: set the maximum learning rate by scaling
                the minimum learning rate. if min_lr = 0.001 and
                lr_range_scaler = 20, then max_lr is set to 0.02.
                this is ignored if max_lr is passed.
            mode: cyclic method. options include:
                [one_cycle, cos_annealing, multi_triangular, exp_range].
            momentum: rate of change for exp_range mode.
                higher numbers slow the the rate of change decrease.
            sample_size: total number of samples in an epoch.
            batch_size: number of samples in a batch.
            step_factor: number of epochs over which to run half a cycle.
                set this to 1 to scale from min to max in one epoch,
                then from max to min in the next.
        """
        super().__init__()

        self.min_lr = min_lr
        self.max_lr = min_lr * lr_range_scaler if max_lr is None else max_lr
        self.mode = mode
        self.momentum = momentum
        self.clr_iterations = 0.0
        self.scale = 1.0
        self.step_factor = step_factor
        if sample_size and batch_size:
            self.step_size = int(round(self.step_factor * sample_size // batch_size))
        else:
            self.step_size = 1000
        self.lrs, self.iterations = [], []

    def lr_scale(self):
        current_cycle = 1 + self.clr_iterations // (2 * self.step_size)
        if self.mode == "cos_annealing":
            cos = 1 + math.cos(math.pi * self.clr_iterations / self.step_size)
            return cos
        elif self.mode == "one_cycle":
            return 1.0
        elif self.mode == "multi_triangular":
            return 1 / (2.0**current_cycle)
        elif self.mode == "exp_range":
            return self.momentum**self.clr_iterations
        else:
            raise ValueError

    def cyclic_lr(self):
        cycle = 1 + self.clr_iterations // (2 * self.step_size)
        x = abs(self.clr_iterations / self.step_size - 2 * cycle + 1)
        if self.mode == "cos_annealing":
            lr = self.min_lr + 0.5 * (self.max_lr - self.min_lr) * self.lr_scale()
        else:
            lr = (
                self.min_lr
                + (self.max_lr - self.min_lr) * max(0.0, (1.0 - x)) * self.lr_scale()
            )

        return lr

    def on_train_begin(self, logs=None):
        super().on_train_begin(logs=logs)
        if self.clr_iterations == 0:
            self.model.optimizer.lr = self.min_lr
        else:
            self.model.optimizer.lr = self.cyclic_lr()

    def on_batch_end(self, batch, logs=None):
        super().on_batch_end(batch, logs=logs)
        self.clr_iterations += 1.0
        self.iterations.append(self.clr_iterations)
        self.model.optimizer.lr = self.cyclic_lr()
        self.lrs.append(self.model.optimizer.lr)
        logs.update({"LR": self.lrs[-1]})

    def on_epoch_end(self, epoch, logs=None):
        self.model.optimizer.lr = self.lrs[-1]

__init__(min_lr=0.0001, max_lr=None, lr_range_scaler=10, mode='one_cycle', momentum=0.8, sample_size=None, batch_size=None, step_factor=1)

Create a cyclic learning rate callback to modify LR by batch.

Reference: L.N. Smith 2017 arxiv.org/pdf/1506.01186.pdf

Parameters:

Name Type Description Default
min_lr float

minimum bounding learning rate.

 0.0001
max_lr float

maximum bounding learning rate.

 None
lr_range_scaler float

set the maximum learning rate by scaling the minimum learning rate. if min_lr = 0.001 and lr_range_scaler = 20, then max_lr is set to 0.02. this is ignored if max_lr is passed.

 10
mode str

cyclic method. options include: [one_cycle, cos_annealing, multi_triangular, exp_range].

 'one_cycle'
momentum float

rate of change for exp_range mode. higher numbers slow the the rate of change decrease.

 0.8
sample_size int

total number of samples in an epoch.

 None
batch_size int

number of samples in a batch.

 None
step_factor int

number of epochs over which to run half a cycle. set this to 1 to scale from min to max in one epoch, then from max to min in the next.

 1
Source code in myco/callbacks.py 
 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
114
115
116
117
118
119
120
121
122
123
def __init__(
    self,
    min_lr: float = 0.0001,
    max_lr: float = None,
    lr_range_scaler: float = 10,
    mode: str = "one_cycle",
    momentum: float = 0.8,
    sample_size: int = None,
    batch_size: int = None,
    step_factor: int = 1,
):
    """Create a cyclic learning rate callback to modify LR by batch.

    Reference: L.N. Smith 2017 https://arxiv.org/pdf/1506.01186.pdf

    Args:
        min_lr: minimum bounding learning rate.
        max_lr: maximum bounding learning rate.
        lr_range_scaler: set the maximum learning rate by scaling
            the minimum learning rate. if min_lr = 0.001 and
            lr_range_scaler = 20, then max_lr is set to 0.02.
            this is ignored if max_lr is passed.
        mode: cyclic method. options include:
            [one_cycle, cos_annealing, multi_triangular, exp_range].
        momentum: rate of change for exp_range mode.
            higher numbers slow the the rate of change decrease.
        sample_size: total number of samples in an epoch.
        batch_size: number of samples in a batch.
        step_factor: number of epochs over which to run half a cycle.
            set this to 1 to scale from min to max in one epoch,
            then from max to min in the next.
    """
    super().__init__()

    self.min_lr = min_lr
    self.max_lr = min_lr * lr_range_scaler if max_lr is None else max_lr
    self.mode = mode
    self.momentum = momentum
    self.clr_iterations = 0.0
    self.scale = 1.0
    self.step_factor = step_factor
    if sample_size and batch_size:
        self.step_size = int(round(self.step_factor * sample_size // batch_size))
    else:
        self.step_size = 1000
    self.lrs, self.iterations = [], []

GarbageCollector

Bases: tf.keras.callbacks.Callback

Clears memory leaks

Source code in myco/callbacks.py 
50
51
52
53
54
55
56
57
58
59
class GarbageCollector(tf.keras.callbacks.Callback):
    """Clears memory leaks"""

    def __init__(self):
        super().__init__()

    def on_epoch_end(self, epoch, logs=None):
        gc.collect()

    K.clear_session()

LrFinder

Bases: tf.keras.callbacks.Callback

This callback uses exponential annealing to find the loss value per lr in a range default range value from start_lr=1e-10 to end_lr=1e1

Source code in myco/callbacks.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
class LrFinder(tf.keras.callbacks.Callback):
    """
    This callback uses exponential annealing to find the loss value per lr in a range
    default range value from start_lr=1e-10 to end_lr=1e1
    """

    def __init__(self, start_lr=1e-9, end_lr=1e-2, n_rounds=1e3):
        super().__init__()
        self.weights, self.batch_loss = None, None
        self.start_lr, self.end_lr, self.n_rounds = start_lr, end_lr, n_rounds
        self.step_up = (end_lr / start_lr) ** (1 / n_rounds)
        self.lrs, self.losses = [], []

    def exp_annealing(self):
        return self.model.optimizer.lr * self.step_up

    def on_train_begin(self, logs=None):
        self.weights = self.model.get_weights()
        self.model.optimizer.lr = self.start_lr

    def on_train_batch_begin(self, batch, logs=None):
        self.model.optimizer.lr = self.exp_annealing()

    def on_train_batch_end(self, batch, logs=None):
        self.losses.append(logs["loss"])
        self.lrs.append(self.model.optimizer.lr.numpy())
        if self.model.optimizer.lr > self.end_lr:
            self.model.stop_training = True
        logs.update({"LR": self.lrs[-1]})
        super().on_train_batch_end(batch, logs)

    def on_train_end(self, logs=None):
        self.model.set_weights(self.weights)

MemoryCallback

Bases: tf.keras.callbacks.Callback

Print memory callbacks

Source code in myco/callbacks.py 
37
38
39
40
41
42
43
44
45
46
47
class MemoryCallback(tf.keras.callbacks.Callback):
    """Print memory callbacks"""

    def __init__(self):
        super().__init__()

    def on_epoch_end(self, epoch, log={}):
        _logger.info(
            f"Memory callback of the epoch: "
            f"{resource.getrusage(resource.RUSAGE_SELF).ru_maxrss}"
        )

MycoTensorBoard

Bases: tf.keras.callbacks.TensorBoard

Extends tensorboard to report additional logging information

Source code in myco/callbacks.py 
62
63
64
65
66
67
68
69
70
71
72
class MycoTensorBoard(tf.keras.callbacks.TensorBoard):
    """Extends tensorboard to report additional logging information"""

    def __init__(self, lr_policy="cyclic", **kwargs):
        self.lr_policy = lr_policy
        super().__init__(**kwargs)

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        logs.update({"LR": self.model.optimizer.lr.numpy()})
        super().on_epoch_end(epoch, logs)

get_required_callbacks(model_path, best_only=True)

Creates list of required model callbacks for model training.

Parameters:

Name Type Description Default
model_path str

Output filepath for model.

 required

Returns:

Type Description
List[tf.keras.callbacks.Callback]

List of model callbacks.
Source code in myco/callbacks.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def get_required_callbacks(
    model_path: str, best_only: bool = True
) -> List[tf.keras.callbacks.Callback]:
    """Creates list of required model callbacks for model training.

    Args:
        model_path: Output filepath for model.

    Returns:
        List of model callbacks.
    """
    # check for nan first no matter what
    callbacks = [
        tf.keras.callbacks.TerminateOnNaN(),
        tf.keras.callbacks.ModelCheckpoint(model_path, save_best_only=best_only),
    ]
    return callbacks