사전 학습 모델을 활용한 이미지 분류

사전 학습 모델을 활용한 이미지 분류

1. 서론

딥러닝 모델을 처음부터 학습하는 것은 많은 데이터와 연산 자원이 필요합니다. 하지만 사전 학습된 모델(Pre-trained Model)을 활용하면 비교적 적은 데이터와 연산 비용으로도 뛰어난 성능을 발휘할 수 있습니다. 이번 글에서는 사전 학습된 모델을 활용하여 이미지 분류를 수행하는 방법에 대해 알아보겠습니다.

2. 사전 학습 모델이란?

사전 학습 모델은 대규모 데이터셋을 이용해 미리 학습된 신경망 모델을 의미합니다. 대표적인 예로 ImageNet 데이터셋을 활용해 학습된 모델들이 있으며, 이를 전이 학습(Transfer Learning) 기법을 이용해 새로운 이미지 분류 문제에 적용할 수 있습니다.

2.1 대표적인 사전 학습 모델

• VGG16, VGG19: 비교적 단순한 구조의 CNN 모델
• ResNet(Residual Network): 딥러닝에서 잔차 연결을 활용한 강력한 모델
• Inception: 다양한 크기의 커널을 활용한 네트워크 구조
• EfficientNet: 연산량 대비 높은 성능을 보이는 모델
• MobileNet: 모바일 기기에서도 사용할 수 있도록 최적화된 경량 모델

3. 전이 학습을 활용한 이미지 분류

사전 학습된 모델을 활용하여 이미지 분류를 수행하는 방법은 크게 두 가지로 나눌 수 있습니다.

1. Feature Extraction (특징 추출): 사전 학습된 모델의 일부 계층을 그대로 사용하여 특징을 추출하고, 새로운 분류기를 추가하여 학습합니다.
2. Fine-tuning (미세 조정): 사전 학습된 모델의 일부 또는 전체를 다시 학습하여 새로운 데이터셋에 적합하도록 조정합니다.

이번 실습에서는 TensorFlow와 Keras를 활용하여 ImageNet으로 학습된 ResNet50 모델을 사용하여 이미지 분류를 수행하는 방법을 살펴보겠습니다.

4. 실습: ResNet50을 활용한 이미지 분류

4.1 필요한 라이브러리 설치 및 불러오기

import tensorflow as tf
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.layers import Dense, Flatten, Dropout
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
import matplotlib.pyplot as plt

4.2 데이터 준비

일반적으로 전이 학습에서는 적은 양의 데이터만으로도 높은 성능을 낼 수 있습니다. 여기서는 TensorFlow에서 제공하는 Cats vs. Dogs 데이터셋을 사용하겠습니다.

import tensorflow_datasets as tfds

dataset, info = tfds.load("cats_vs_dogs", as_supervised=True, with_info=True)
train_dataset = dataset['train']

# 데이터 샘플 시각화
plt.figure(figsize=(10,5))
for i, (image, label) in enumerate(train_dataset.take(5)):
    plt.subplot(1, 5, i+1)
    plt.imshow(image)
    plt.axis('off')
plt.show()

4.3 데이터 전처리

ResNet50 모델은 특정한 입력 크기를 요구하므로, 이를 맞추고 데이터 증강을 수행합니다.

def preprocess(image, label):
    image = tf.image.resize(image, (224, 224)) / 255.0
    return image, label

batch_size = 32
train_data = train_dataset.map(preprocess).batch(batch_size).prefetch(tf.data.AUTOTUNE)

4.4 ResNet50 모델 불러오기

사전 학습된 ResNet50 모델을 기반으로 새로운 분류기를 추가합니다.

base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
base_model.trainable = False  # Feature Extraction 방식 적용

x = Flatten()(base_model.output)
x = Dense(256, activation='relu')(x)
x = Dropout(0.5)(x)
output = Dense(1, activation='sigmoid')(x)

model = Model(inputs=base_model.input, outputs=output)

4.5 모델 컴파일 및 학습

모델을 컴파일하고 학습을 진행합니다.

model.compile(optimizer=Adam(learning_rate=0.001), loss='binary_crossentropy', metrics=['accuracy'])
model.fit(train_data, epochs=5)

4.6 모델 성능 평가

테스트 이미지를 활용하여 모델이 정상적으로 동작하는지 확인합니다.

import numpy as np

def predict_image(image):
    image = tf.image.resize(image, (224, 224)) / 255.0
    image = tf.expand_dims(image, axis=0)  # 배치 차원 추가
    prediction = model.predict(image)
    return 'Dog' if prediction > 0.5 else 'Cat'

plt.figure(figsize=(5, 5))
for image, label in train_dataset.take(1):
    plt.imshow(image)
    plt.axis('off')
    plt.title(predict_image(image))
plt.show()

5. 미세 조정 (Fine-tuning)

이제 Feature Extraction 방식이 아닌, 일부 계층을 학습 가능한 상태로 두고 미세 조정을 진행해 보겠습니다.

base_model.trainable = True  # 사전 학습된 모델의 학습 허용

# 특정 계층 이후부터 학습을 진행하도록 설정
for layer in base_model.layers[:-30]:
    layer.trainable = False

model.compile(optimizer=Adam(learning_rate=0.0001), loss='binary_crossentropy', metrics=['accuracy'])
model.fit(train_data, epochs=5)

6. 결론

이 글에서는 사전 학습된 ResNet50 모델을 활용하여 이미지 분류를 수행하는 방법을 살펴보았습니다. Feature Extraction 방식과 Fine-tuning 방식을 비교하며 학습 방법을 적용해 보았습니다. 전이 학습 기법을 활용하면 제한된 데이터와 자원으로도 강력한 성능을 얻을 수 있습니다.