Is there a way in Keras to immediately stop training?

Viewed 1671

I am writing a custom early stopping callback for my tf.keras training. For that I can set the variable self.model.stop_training = True in one of the callback functions, like for example on_epoch_end(). However, Keras stops the training only when the current epoch is done, even if I set this variable within the training of one epoch, for example in on_batch_end().

Hence my question: Is there a way in Keras to stop training immediately, even within the progress of the current epoch?

2 Answers

In keras you use EarlyStopping to stop when a monitored quantity has stopped improving. From your question it is not clear what the condition is you want to stop. If you just want to monitor a value like in EarlyStopping but only want to stop after a batch, if the value is not improving, you can just rewrite the EarlyStopping class and implement the logic in on_batch_end instead of on_epoch_end:

class EarlyBatchStopping(Callback):


    def __init__(self,
                 monitor='val_loss',
                 min_delta=0,
                 patience=0,
                 verbose=0,
                 mode='auto',
                 baseline=None,
                 restore_best_weights=False):
        super(EarlyStopping, self).__init__()

        self.monitor = monitor
        self.baseline = baseline
        self.patience = patience
        self.verbose = verbose
        self.min_delta = min_delta
        self.wait = 0
        self.stopped_epoch = 0
        self.restore_best_weights = restore_best_weights
        self.best_weights = None

        if mode not in ['auto', 'min', 'max']:
            warnings.warn('EarlyStopping mode %s is unknown, '
                          'fallback to auto mode.' % mode,
                          RuntimeWarning)
            mode = 'auto'

        if mode == 'min':
            self.monitor_op = np.less
        elif mode == 'max':
            self.monitor_op = np.greater
        else:
            if 'acc' in self.monitor:
                self.monitor_op = np.greater
            else:
                self.monitor_op = np.less

        if self.monitor_op == np.greater:
            self.min_delta *= 1
        else:
            self.min_delta *= -1

    def on_train_begin(self, logs=None):
        # Allow instances to be re-used
        self.wait = 0
        self.stopped_epoch = 0
        if self.baseline is not None:
            self.best = self.baseline
        else:
            self.best = np.Inf if self.monitor_op == np.less else -np.Inf

    def on_batch_end(self, epoch, logs=None):
        current = self.get_monitor_value(logs)
        if current is None:
            return

        if self.monitor_op(current - self.min_delta, self.best):
            self.best = current
            self.wait = 0
            if self.restore_best_weights:
                self.best_weights = self.model.get_weights()
        else:
            self.wait += 1
            if self.wait >= self.patience:
                self.stopped_epoch = epoch
                self.model.stop_training = True
                if self.restore_best_weights:
                    if self.verbose > 0:
                        print('Restoring model weights from the end of '
                              'the best epoch')
                    self.model.set_weights(self.best_weights)

    def on_train_end(self, logs=None):
        if self.stopped_epoch > 0 and self.verbose > 0:
            print('Epoch %05d: early stopping' % (self.stopped_epoch + 1))

    def get_monitor_value(self, logs):
        monitor_value = logs.get(self.monitor)
        if monitor_value is None:
            warnings.warn(
                'Early stopping conditioned on metric `%s` '
                'which is not available. Available metrics are: %s' %
                (self.monitor, ','.join(list(logs.keys()))), RuntimeWarning
            )
        return monitor_value

If you have another logic, you can use on_batch_end and set self.model.stop_training = True based on your logic, but I think you got the idea.

You can use model.stop_training parameter to stop the training.

For example, if we want to stop the training at 2nd epochs 3rd batch then you can do something like below.

import keras
from keras.models import Sequential
from keras.layers import Dense
from keras.optimizers import SGD
import numpy as np
import pandas as pd

class My_Callback(keras.callbacks.Callback):
    def on_epoch_begin(self, epoch, logs={}):
      self.epoch = epoch

    def on_batch_end(self, batch, logs={}):
        if self.epoch == 1 and batch == 3:
          print (f"\nStopping at Epoch {self.epoch}, Batch {batch}")
          self.model.stop_training = True


X_train = np.random.random((100, 3))
y_train = pd.get_dummies(np.argmax(X_train[:, :3], axis=1)).values

clf = Sequential()
clf.add(Dense(9, activation='relu', input_dim=3))
clf.add(Dense(3, activation='softmax'))
clf.compile(loss='categorical_crossentropy', optimizer=SGD())

clf.fit(X_train, y_train, epochs=10, batch_size=16, callbacks=[My_Callback()])

Output:

Epoch 1/10
100/100 [==============================] - 0s 337us/step - loss: 1.0860
Epoch 2/10
 16/100 [===>..........................] - ETA: 0s - loss: 1.0830
Stopping at Epoch 1, Batch 3
<keras.callbacks.callbacks.History at 0x7ff2e3eeee10>
Related