Source code for flax.training.early_stopping

# Copyright 2022 The Flax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Early stopping."""

import math
from flax import struct


[docs]class EarlyStopping(struct.PyTreeNode): """Early stopping to avoid overfitting during training. The following example stops training early if the difference between losses recorded in the current epoch and previous epoch is less than 1e-3 consecutively for 2 times:: early_stop = EarlyStopping(min_delta=1e-3, patience=2) for epoch in range(1, num_epochs+1): rng, input_rng = jax.random.split(rng) optimizer, train_metrics = train_epoch( optimizer, train_ds, config.batch_size, epoch, input_rng) _, early_stop = early_stop.update(train_metrics['loss']) if early_stop.should_stop: print('Met early stopping criteria, breaking...') break Attributes: min_delta: Minimum delta between updates to be considered an improvement. patience: Number of steps of no improvement before stopping. best_metric: Current best metric value. patience_count: Number of steps since last improving update. should_stop: Whether the training loop should stop to avoid overfitting. """ min_delta: float = 0 patience: int = 0 best_metric: float = float('inf') patience_count: int = 0 should_stop: bool = False def reset(self): return self.replace(best_metric=float('inf'), patience_count=0, should_stop=False)
[docs] def update(self, metric): """Update the state based on metric. Returns: A pair (has_improved, early_stop), where `has_improved` is True when there was an improvement greater than `min_delta` from the previous `best_metric` and `early_stop` is the updated `EarlyStop` object. """ if math.isinf(self.best_metric) or self.best_metric - metric > self.min_delta: return True, self.replace(best_metric=metric, patience_count=0) else: should_stop = self.patience_count >= self.patience or self.should_stop return False, self.replace(patience_count=self.patience_count + 1, should_stop=should_stop)