개발일기

이 글에서는 사전 훈련된 네트워크에서 Transfer Learning(전이 학습)을 사용하여 이미지를 분류하는 방법을 다룬다.

IMG_SHAPE = (128, 128, 3)
base_model = tf.keras.applications.MobileNetV2(input_shape = IMG_SHAPE, include_top = False)

이미 만들어져서 학습이 완료된 모델을 가져오되, 헤드모델은 제외한 베이스모델만 가져온다.

base_model.trainable = False

우리가 가져온 모델의 베이스모델 부분은특징을 잡아내는 역활을 하므로, 학습이 되지 않도록 한다.

head_model = base_model.output
head_model = tf.keras.layers.GlobalAveragePooling2D()(head_model)
head_model = tf.keras.layers.Dense(128, 'relu')(head_model)
head_model = tf.keras.layers.Dense(1, 'sigmoid')(head_model)

head_model을 정의해준다.

model = tf.keras.models.Model(inputs = base_model.input, outputs = head_model)
model.compile(tf.keras.optimizers.RMSprop(0.0001), 'binary_crossentropy', ['accuracy'])

베이스 모델과 헤드모델을 합쳐서, 하나의 모델으로 변수로 만들어 준다.

이후 model.fit을 하면 된다.