为何Keras自定义on_train_batch_end回调无法中途停止训练？求解决方案

Hey there! I get it—waiting out a full 1.5-hour epoch when your model already hit the target accuracy minutes ago is such a waste of time. Let's figure out why your callback isn't stopping the training and get it working right.

First, Let's Fix Your Callback #

Your core idea is correct, but a few small tweaks will make sure the stop signal gets picked up properly. Here's the adjusted version:

import tensorflow as tf

acc_thresh = 0.965
class myCallback(tf.keras.callbacks.Callback):
    def on_train_batch_end(self, batch, logs=None):
        # Guard against empty logs to avoid KeyErrors
        logs = logs or {}
        current_acc = logs.get('accuracy')
        
        # Only trigger if we have a valid accuracy reading that meets the threshold
        if current_acc is not None and current_acc > acc_thresh:
            print(f"\nWe've reached {acc_thresh*100:.2f}% accuracy—stopping training immediately!")
            self.model.stop_training = True
            # Optional: Clear the session to free up resources right away
            tf.keras.backend.clear_session()

Key Things to Check #

• Double-Check the Accuracy Key in Logs
  Depending on how you compiled your model, the accuracy metric might have a different key in the logs. For example:

  • If you used model.compile(..., metrics=['accuracy']), logs.get('accuracy') is correct.
  • If you used sparse categorical crossentropy (for integer labels), the key might be sparse_categorical_accuracy.
    Always match the key to the metrics you specified during compilation!

• Understand When the Stop Happens
  Setting self.model.stop_training = True doesn't halt the current batch mid-execution—it tells TensorFlow to skip all subsequent batches in the epoch. If you see the print message but training keeps going for a bit, it's just finishing up the current batch (which makes sense if your batches are large).

• Make Sure the Callback is Registered
  Don't forget to pass your callback to model.fit()—it sounds obvious, but it's an easy miss:

  callbacks = [myCallback()]
  model.fit(
      x_train, y_train,
      epochs=100,
      batch_size=32,
      callbacks=callbacks
  )
  

• Watch for Distributed Training Gotchas
  If you're using a distributed strategy like MirroredStrategy, you'll need to ensure the stop signal is propagated across all replicas. In most cases, setting self.model.stop_training = True on the main replica will work, but if not, add logic to check if you're on the main replica before triggering the stop.

Test with Verbose Logging #

If you're still having trouble, add some extra logging to verify what's happening each batch:

def on_train_batch_end(self, batch, logs=None):
    logs = logs or {}
    current_acc = logs.get('accuracy')
    print(f"Batch {batch} | Current Accuracy: {current_acc:.4f}")
    
    if current_acc is not None and current_acc > acc_thresh:
        print(f"\nThreshold hit! Stopping after batch {batch}.")
        self.model.stop_training = True

This will let you see exactly which batch triggers the stop, and confirm that training stops after that batch completes.

内容的提问来源于stack exchange，提问作者Timothy Alex John