개발공부

텐서플로우의 콜백 클래스를 이용해서, 원하는 조건이 되면 학습을 멈추게 하는 코드 본문

Python/Deep Learning

텐서플로우의 콜백 클래스를 이용해서, 원하는 조건이 되면 학습을 멈추게 하는 코드

mscha 2022. 6. 13. 17:58

내가 정한 수치에 도달하면, 학습을 멈추게 하는 방법

텐서 플로우의 콜백 클래스를 상속해서 만든다.

함수 on_epoch_end 함수 안에, 에포크가 끝날 때마다 하고 싶은 작업을, 코딩을 해주면 된다.

# 모델생성 함수
def build_model() :
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(128, 'relu'))
    # 3개 이상의 분류 문제 Output Layer activation func -> softmax
    model.add(tf.keras.layers.Dense(10, 'softmax'))
    model.compile(optimizer = 'adam', loss = 'sparse_categorical_crossentropy', metrics = ['accuracy'])
    return model
# val_accuracy > 0.87 일 때 학습 종료
class myCallback(tf.keras.callbacks.Callback) :
    def on_epoch_end(self, epoch, logs={}) :
        if logs.get('val_accuracy') > 0.87 :
            print('\n밸리데이션 정확도가 87%를 넘으므로, 학습을 종료합니다.')
            self.model.stop_training = True
# 변수에 저장
my_callback = myCallback()
# 모델 생성
model9 = build_model()
# 모델 학습
epoch_history = model9.fit(training_images, training_labels, epochs = 1000, validation_split = 0.2, 
						callbacks = [my_callback])