Skip weights of one epoch in Tensorflow Keras

Viewed 269

In my machine learning task I have the problem, that in some rare cases (epochs) the optimiser sees a bad training set and the weights and biases get completely messed up after that epoch, so I would like to discard this epoch.

I wrote a small callback function like in this example

from keras.callbacks import Callback

class skip_bad_epoch(Callback):
    """Skip epoch when the loss significantly higher than the previous epoch.
  Arguments:
      factor: Factor by which the loss has to be higher than the previous   """

    def __init__(self, factor=2):
        super(skip_bad_epoch, self).__init__()
        self.factor = factor
        # best_weights to store the weights at which the minimum loss occurs.
        self.last_weights = None

    def on_train_begin(self, logs=None):
        # Initialize the last as infinity.
        self.last = np.Inf

    def on_epoch_end(self, epoch, logs=None):
        current = logs.get("loss") # Use train loss here because that really freaks out
        if (current/self.last)<self.factor: # If the current loss is not significantly higher than the previous
            self.last = current
            # Record the best weights if current results is better (less loss).
            self.last_weights = self.model.get_weights()
        else:
            print(f"Restoring model weights from the previous epoch. Last loss: {self.last} this loss {current}")
            self.model.set_weights(self.last_weights)

And called it during training like so:

callbacks_list = [skip_bad_epoch(factor=1.5)]
model.fit(
    x, y,
    validation_split=0.1,
    callbacks=callbacks_list
    )

Results: In this example epoch #10 was bad, and the callback tries so skip it. However the loss in the upcoming epoch (here: 11) is always significantly worse than the ones before the skipped one. So it seems like the weights are still adjusted and resetting them with self.model.set_weights(self.last_weights) is not enough? Maybe this does not include the biases?

Epoch 00007: val_loss did not improve from 0.02599
50/50 [==============================] - 6s 117ms/step - loss: 0.0258 - val_loss: 0.0262

Epoch 00008: val_loss improved from 0.02599 to 0.02570, saving model
50/50 [==============================] - 6s 121ms/step - loss: 0.0260 - val_loss: 0.0257

Epoch 00009: val_loss did not improve from 0.02570
50/50 [==============================] - 6s 120ms/step - loss: 0.0259 - val_loss: 0.0475

Epoch 00010: val_loss did not improve from 0.02570
Restoring model weights from the previous epoch. Last loss: 0.0259 this loss 0.0444
50/50 [==============================] - 6s 119ms/step - loss: 0.0444 - val_loss: 0.0331

Epoch 00011: val_loss did not improve from 0.02570
50/50 [==============================] - 6s 121ms/step - loss: 0.0292 - val_loss: 0.0264

Epoch 00012: val_loss did not improve from 0.02570
50/50 [==============================] - 6s 129ms/step - loss: 0.0261 - val_loss: 0.0259

enter image description here

Using: Tensorflow 2.3.1 Keras 2.4.3 Python 3.8.6

1 Answers

I built a callback doing exactly what you wish to do, In addition it adjust the learning as well. First it adjusts the learning rate by monitoring training accuracy. Once the training accuracy reaches a threshold level, say .95 then the callback adjust the learning rate based on validation loss. Validation data when the training accuracy is low is pretty much irrelevant. Like you I thought if you get a significant decrease in performance of the quantity being monitored at the end of an epoch you should do two things, One is to lower the learning rate for sure. The other is to restore the weights from the previous epoch. Load and save weights includes loading and saving the biases by the way. Logic is that the metric being monitored is a surface in N space where N is the number of trainable parameters. If you get a degradation in performance you have moved into a location in N space that is not as good as the one you were in in the prior epoch. Why stay there. So load the weights from the previous epoch and start training again from that better place but with a reduced learning rate. I think where you saw the loss increase in epoch 11 was because your learning rate was to large which is probably also what caused the large increase in epoch 10. When I reduce the lr I multiply the lr by a fixed factor like .5 then I further multiply it by the ratio of lowest v_loss/v_loss current epoch. I have a parameter called dwell which I can set True or False to either enable or disable the reloading of the weights. I have noticed some improvement in faster model convergence(not much) in using the dwell idea but I have not experimented that much. One danger I see is using this dwelling might quickly get you stuck in a local minimum in the sense you do not "explore" N space as much. So try the idea of lowering the lr and see if it works better. You can get the current learning rate and set the new learning with the code shown below:

current_lr=float(tf.keras.backend.get_value(self.model.optimizer.lr))
new_lr= whatever calculation you want to use
tf.keras.backend.set_value(self.model.optimizer.lr, new_lr)
Related