Tensorflow - Callback 클래스를 이용해서, 원하는 조건이 되면 학습을 멈추기

accuracy나 loss 등이, 내가 원하는 특정 값이 되면 자동으로 학습을 멈추게 하고 싶다.

Tensorflow에서는 특정값에 도달하면 학습을 멈추게 할 수 있는 콜백 기능을 제공한다.

 

tf.keras.callbacks.Callback을 상속해서 사용한다.

 

import tensorflow as tf
from tensorflow import keras
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import Flatten

class myCallback(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs={}):
    if logs['val_accuracy'] > 0.88:
      print('\n내가 정한 정확도에 도달했으니, 학습을 멈춘다')
      self.model.stop_training = True

my_cb = myCallback()

def build_model():
  model = Sequential()
  model.add(Flatten())
  model.add(Dense(128, 'relu'))
  model.add(Dense(64, 'relu')) 
  model.add(Dense(10, 'softmax'))
  model.compile('adam', 'sparse_categorical_crossentropy', ['accuracy'])
  return model
  
model = build_model()

epoch_history = model.fit(X_train, y_train, epochs=30, validation_split=0.2,
                          callbacks= [my_cb])