전이 학습 실전 - Alien vs Predator · 콘크리트 균열 탐지, AlexNet · VGG19 활용

소규모 데이터셋에서 사전학습 모델 100% 활용하기

대규모 데이터셋으로 사전 학습된 모델을 가져와, 새로운 작은 데이터셋에 전이 학습(Transfer Learning) 으로 적용하면 적은 데이터와 자원으로도 높은 성능을 낼 수 있다. 이 글에서는 영화 캐릭터 분류용 Alien vs Predator 데이터셋과 콘크리트 균열 탐지용 Surface Crack Detection 데이터셋에 대해, 각각 AlexNet과 VGG19를 활용해 이진 분류 모델을 학습시키고, Optimizer에 따른 성능 비교까지 진행한다.

(Alien vs. Predator 데이터셋, Surface Crack Detection 데이터셋을 기반으로 재구성)

---

1. Alien vs Predator 데이터셋

Alien vs Predator 데이터셋은 영화 속 캐릭터인 에일리언(Alien) 과 프레데터(Predator) 이미지를 모아 둔 소규모 이미지 분류용 데이터셋이다. 두 클래스를 구분하는 CNN을 학습시키기에 좋은 예제다.

1-1. Kaggle에서 다운로드 (Colab 기준)

from google.colab import files
files.upload()  # kaggle.json 업로드 창

!mkdir -p ~/.kaggle
!cp kaggle.json ~/.kaggle/
!chmod 600 ~/.kaggle/kaggle.json

!kaggle datasets download -d pmigdal/alien-vs-predator-images
!unzip -q alien-vs-predator-images.zip

데이터셋 구조는 대략 data/train/alien, data/train/predator, data/validation/... 형태로 맞춰 둔다.

---

2. AlexNet 전이 학습으로 Alien vs Predator 분류

사전 학습된 AlexNet(ImageNet 기반)을 불러와, 마지막 분류기를 Alien/Predator 이진 분류에 맞게 교체한 뒤 학습한다.

import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(device)

2-1. 전처리 및 DataLoader

입력 크기는 PyTorch AlexNet에 맞춰 224×224로 맞추고, 학습용에는 랜덤 Affine·좌우 반전을 사용해 데이터 증강을 준다.

data_transforms = {
    'train': transforms.Compose([
        transforms.Resize((224, 224)),     # 파이토치 AlexNet은 224×224 입력
        transforms.RandomAffine(0, shear=10, scale=(0.8, 1.2)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
    ]),
    'validation': transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])
}

레이블을 1차원 float 텐서로 바꿔 BCELoss와 함께 사용할 수 있게 한다.

def target_transforms(target):
    return torch.FloatTensor([target])

image_datasets = {
    'train': datasets.ImageFolder('data/train', data_transforms['train'],
                                  target_transform=target_transforms),
    'validation': datasets.ImageFolder('data/validation', data_transforms['validation'],
                                       target_transform=target_transforms),
}

dataloaders = {
    'train': DataLoader(image_datasets['train'], batch_size=32, shuffle=True),
    'validation': DataLoader(image_datasets['validation'], batch_size=32, shuffle=False),
}

len(image_datasets['train']), len(image_datasets['validation'])

샘플 이미지를 확인해 본다.

imgs, labels = next(iter(dataloaders['train']))

fig, axes = plt.subplots(4, 8, figsize=(16, 8))
for ax, img, label in zip(axes.flatten(), imgs, labels):
    ax.imshow(img.permute(1, 2, 0))  # (3, 224, 224) → (224, 224, 3)
    ax.set_title(label.item())
    ax.axis('off')

2-2. 전이 학습(Transfer Learning) 개념

전이 학습은 대규모 데이터셋으로 사전 학습된 모델의 가중치를 가져와, 새로운 데이터셋에 일부 레이어만 학습시키는 방법이다. 일반적인 Edge·Texture 등은 그대로 두고(초기 레이어 동결), 마지막 분류 헤드만 현재 문제에 맞게 학습시키면 적은 데이터로도 좋은 성능을 낼 수 있다.

※ ImageNet

ImageNet은 약 1,400만 장, 22,000개 이상 카테고리로 구성된 대규모 이미지 데이터셋이다. 그중 ILSVRC 버전은 약 1,000클래스와 120만 장 이미지를 포함해, ResNet·VGG·Inception 등 수많은 모델의 사전 학습에 쓰이며 전이 학습의 표준 벤치마크로 활용된다.

---

3. AlexNet 사전 학습 모델 불러오기 및 Freeze

model = models.alexnet(weights='IMAGENET1K_V1').to(device)
print(model)

모든 파라미터를 동결(freeze) 하고, 마지막 분류기만 새로 학습한다.

for param in model.parameters():
    param.requires_grad = False

※ Model Freezing

사전 학습된 모델에서 일반적인 특징(에지, 패턴 등) 을 학습한 합성곱 계층은 고정해 두고, 새로운 데이터셋에 특화된 부분(주로 마지막 FC 레이어)만 학습하면, 학습해야 할 파라미터 수가 줄어들어 속도·안정성·과적합 방지 측면에서 유리하다. 필요하면 나중에 일부 레이어의 동결을 풀어 미세조정(Fine-tuning)을 할 수 있다.

---

4. AlexNet 분류기 교체 및 학습

Alien vs Predator는 이진 분류이므로, classifier를 1차원 출력 + Sigmoid로 교체한다.

model.classifier = nn.Sequential(
    nn.Linear(256 * 6 * 6, 128),
    nn.ReLU(),
    nn.Linear(128, 1),
    nn.Sigmoid()
).to(device)

print(model)

학습 루프는 train/validation 두 phase를 번갈아 돌면서 손실과 정확도를 출력한다.

optimizer = optim.Adam(model.classifier.parameters(), lr=0.001)
epochs = 10

for epoch in range(epochs):
    for phase in ['train', 'validation']:
        if phase == 'train':
            model.train()
        else:
            model.eval()

        sum_losses = 0
        sum_accs = 0

        for x_batch, y_batch in dataloaders[phase]:
            x_batch = x_batch.to(device)
            y_batch = y_batch.to(device)

            y_pred = model(x_batch)
            loss = nn.BCELoss()(y_pred, y_batch)

            if phase == 'train':
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

            sum_losses = sum_losses + loss
            y_bool = (y_pred >= 0.5).float()
            acc = (y_batch == y_bool).float().sum() / len(y_batch) * 100
            sum_accs = sum_accs + acc

        avg_loss = sum_losses / len(dataloaders[phase])
        avg_acc = sum_accs / len(dataloaders[phase])
        print(f'{phase:10s}: Epoch {epoch+1:4d}/{epochs} Loss: {avg_loss:.4f} Accuracy: {avg_acc:.2f}%')

4-1. 개별 이미지 추론

from PIL import Image

img1 = Image.open('/content/data/validation/alien/21.jpg')
img2 = Image.open('/content/data/validation/predator/30.jpg')

fig, axes = plt.subplots(1, 2, figsize=(12, 6))
axes[0].imshow(img1); axes[0].axis('off')
axes[1].imshow(img2); axes[1].axis('off')
plt.show()

img1_input = data_transforms['validation'](img1)
img2_input = data_transforms['validation'](img2)
print(img1_input.shape, img2_input.shape)

test_batch = torch.stack([img1_input, img2_input]).to(device)
y_pred = model(test_batch)
print(y_pred)

fig, axes = plt.subplots(1, 2, figsize=(12, 6))
axes[0].set_title(f'{(1-y_pred[0, 0])*100:.2f}% Alien, {(y_pred[0, 0])*100:.2f}% Predator')
axes[0].imshow(img1); axes[0].axis('off')

axes[1].set_title(f'{(1-y_pred[1, 0])*100:.2f}% Alien, {(y_pred[1, 0])*100:.2f}% Predator')
axes[1].imshow(img2); axes[1].axis('off')
plt.show()

---

5. Surface Crack Detection 데이터셋

Surface Crack Detection 데이터셋은 콘크리트 표면 이미지에서 균열(crack) 이 있는지 여부를 이진 분류하는 데 쓰이는 데이터셋이다. 균열 유무에 따라 이미지가 positive/negative 로 나뉜다.

5-1. Kaggle 다운로드 및 압축 해제

from google.colab import files
files.upload()  # kaggle.json 업로드

!mkdir -p ~/.kaggle
!cp kaggle.json ~/.kaggle/
!chmod 600 ~/.kaggle/kaggle.json

!kaggle datasets download arunrk7/surface-crack-detection
!unzip -q surface-crack-detection.zip

---

6. 데이터 분할 및 샘플 시각화

데이터를 train/val/test 와 negative/positive 구조로 나누고, 일부 샘플을 확인한다.

import os
import shutil
import random
import glob
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from torchvision.datasets import ImageFolder
from torchvision.models import vgg19, VGG19_Weights
from torchvision import transforms, models
from torch.utils.data import DataLoader, Subset
from PIL import Image

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(device)

dataset_structure = ['train/negative', 'train/positive',
                     'val/negative', 'val/positive',
                     'test/negative', 'test/positive']

for folder in dataset_structure:
    os.makedirs(folder, exist_ok=True)

categories = ['Negative', 'Positive']

for category in categories:
    files = os.listdir(category)
    random.shuffle(files)
    num_files = len(files)
    train_split = int(num_files * 0.6)
    val_split = int(num_files * 0.2)

    train_files = files[:train_split]
    val_files = files[train_split:train_split + val_split]
    test_files = files[train_split + val_split:]

    target_category = category.lower()

    for file in train_files:
        shutil.copy(os.path.join(category, file), f'train/{target_category}/')
    for file in val_files:
        shutil.copy(os.path.join(category, file), f'val/{target_category}/')
    for file in test_files:
        shutil.copy(os.path.join(category, file), f'test/{target_category}/')

print(\"데이터 분할 및 복사 완료!\")

훈련 데이터에서 positive/negative 각 4장씩 시각화한다.

train_positive_dir = 'train/positive'
train_negative_dir = 'train/negative'

positive_files = random.sample(os.listdir(train_positive_dir), 4)
negative_files = random.sample(os.listdir(train_negative_dir), 4)

positive_paths = [os.path.join(train_positive_dir, f) for f in positive_files]
negative_paths = [os.path.join(train_negative_dir, f) for f in negative_files]

fig, axes = plt.subplots(2, 4, figsize=(12, 6))
fig.suptitle('Train Dataset Samples', fontsize=16)

for i, file_path in enumerate(positive_paths):
    image = Image.open(file_path)
    axes[0, i].imshow(image); axes[0, i].axis('off'); axes[0, i].set_title('Positive')

for i, file_path in enumerate(negative_paths):
    image = Image.open(file_path)
    axes[1, i].imshow(image); axes[1, i].axis('off'); axes[1, i].set_title('Negative')
plt.show()

---

7. 전처리 및 Subset 구성

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

train_dataset = ImageFolder('/content/train', transform=transform)
val_dataset = ImageFolder('/content/val', transform=transform)

print(len(train_dataset), len(val_dataset))

# 예: 전체의 10%만 사용
subset_ratio = 0.1
total_train_size = len(train_dataset)
total_val_size = len(val_dataset)

train_subset_size = int(total_train_size * subset_ratio)
val_subset_size = int(total_val_size * subset_ratio)

train_indices = np.random.choice(total_train_size, train_subset_size, replace=False)
val_indices = np.random.choice(total_val_size, val_subset_size, replace=False)

train_dataset = Subset(train_dataset, train_indices)
val_dataset = Subset(val_dataset, val_indices)

print(f\"Train dataset size after reduction: {len(train_dataset)}\")
print(f\"Validation dataset size after reduction: {len(val_dataset)}\")

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=2)

---

8. VGG19 전이 학습 (균열 탐지)

8-1. VGG19 불러오기 및 분류기 교체

model = vgg19(weights=VGG19_Weights.IMAGENET1K_V1)
print(model)

for param in model.parameters():
    param.requires_grad = False

model.classifier[6] = nn.Linear(4096, 2)
model.classifier[6].requires_grad = True

criterion = nn.CrossEntropyLoss()

※ Adam vs RAdam

• Adam: 가장 널리 쓰이는 옵티마이저로, 대부분의 모델에서 무난하게 잘 동작.\n- RAdam: 학습 초반 손실이 불안정하거나 배치 크기가 작을 때 Adam보다 안정적인 경우가 많다.\n\n일반적으로 Adam으로 시작하고, 학습이 요동치면 RAdam을 고려하는 전략을 쓸 수 있다.

---

9. Optimizer별 성능 비교 함수

import torch.optim as optim

def train_model(optimizer_name, model, train_loader, val_loader, criterion, num_epochs=20):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    if optimizer_name == 'SGD':
        optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    elif optimizer_name == 'Adam':
        optimizer = optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))
    elif optimizer_name == 'RAdam':
        optimizer = optim.RAdam(model.parameters(), lr=0.001, betas=(0.9, 0.999))
    else:
        raise ValueError('알 수 없는 최적화 알고리즘')

    train_losses = []
    val_losses = []
    val_accuracies = []

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0

        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        train_loss = running_loss / len(train_loader)
        train_losses.append(train_loss)

        val_loss = 0.0
        model.eval()
        correct = 0
        total = 0

        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                loss = criterion(outputs, labels)
                val_loss += loss.item()

        val_loss /= len(val_loader)
        val_losses.append(val_loss)
        val_accuracy = 100 * correct / total
        val_accuracies.append(val_accuracy)

        print(f'[{optimizer_name}] Epoch {epoch + 1}, Train Loss: {train_loss: .6f}, '
              f'Val Loss: {val_loss: .6f}, Validation Accuracy: {val_accuracy: .2f}%', flush=True)

    return train_losses, val_losses, val_accuracies

9-1. SGD / Adam / RAdam 비교 학습

train_losses_SGD, val_losses_SGD, val_accuracies_SGD = train_model('SGD', model, train_loader, val_loader, criterion)

# Adam용으로 모델 초기화
model = models.vgg19(pretrained=True)
for param in model.parameters():
    param.requires_grad = False
model.classifier[6] = nn.Linear(4096, 2)
model.classifier[6].requires_grad = True
train_losses_Adam, val_losses_Adam, val_accuracies_Adam = train_model('Adam', model, train_loader, val_loader, criterion)

# RAdam용으로 모델 초기화
model = vgg19(weights=VGG19_Weights.IMAGENET1K_V1)
for param in model.parameters():
    param.requires_grad = False
model.classifier[6] = nn.Linear(4096, 2)
model.classifier[6].requires_grad = True
train_losses_RAdam, val_losses_RAdam, val_accuracies_RAdam = train_model('RAdam', model, train_loader, val_loader, criterion)

---

10. Optimizer별 학습 곡선 시각화

plt.figure(figsize=(15, 10))

# 학습 손실
plt.subplot(3, 1, 1)
plt.plot(train_losses_SGD, label='SGD')
plt.plot(train_losses_Adam, label='Adam')
plt.plot(train_losses_RAdam, label='RAdam')
plt.xlabel('Epoch'); plt.ylabel('Loss')
plt.title('Training Loss Over Epochs')
plt.legend()

# 검증 손실
plt.subplot(3, 1, 2)
plt.plot(val_losses_SGD, label='SGD')
plt.plot(val_losses_Adam, label='Adam')
plt.plot(val_losses_RAdam, label='RAdam')
plt.xlabel('Epoch'); plt.ylabel('Loss')
plt.title('Validation Loss Over Epochs')
plt.legend()

# 검증 정확도
plt.subplot(3, 1, 3)
plt.plot(val_accuracies_SGD, label='SGD', color='blue')
plt.plot(val_accuracies_Adam, label='Adam', color='green')
plt.plot(val_accuracies_RAdam, label='RAdam', color='orange')
plt.xlabel('Epoch'); plt.ylabel('Accuracy (%)')
plt.title('Validation Accuracy Over Epochs')
plt.legend()

plt.tight_layout()
plt.show()

---

11. 테스트 셋에서 예측 결과 시각화

마지막으로 균열/비균열 이미지를 몇 장 골라 VGG19 모델이 어떻게 예측하는지 본다.

def load_and_transform_image(image_path, transform):
    image = Image.open(image_path).convert('RGB')
    return transform(image).unsqueeze(0)

class_folders = {
    'crack': '/content/test/positive',
    'normal': '/content/test/negative',
}

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

plt.figure(figsize=(20, 8))
counter = 1

for class_name, folder_path in class_folders.items():
    image_paths = glob.glob(os.path.join(folder_path, '*'))
    selected_paths = image_paths[:5]

    for image_path in selected_paths:
        image = load_and_transform_image(image_path, transform).to(device)

        model.eval()
        with torch.no_grad():
            outputs = model(image)
            _, predicted = torch.max(outputs, 1)

        prediction = 'normal' if predicted.item() == 0 else 'crack'

        plt.subplot(2, 5, counter)
        plt.imshow(Image.open(image_path))
        plt.title(f'True: {class_name}, Pred: {prediction}')
        plt.axis('off')
        counter += 1

plt.tight_layout()
plt.show()

---

마치며

• Alien vs Predator 데이터셋에는 AlexNet 사전학습 모델을, Surface Crack Detection 에는 VGG19를 전이 학습으로 적용해 소규모 데이터셋에서 효율적으로 학습하는 방법을 살펴봤다.
• 전이 학습에서는 사전학습 가중치 동결 + 분류기 교체가 기본 패턴이며, 상황에 따라 일부 레이어만 미세조정(Fine-tuning)할 수 있다.
• Optimizer로는 보통 Adam 을 기본값으로 쓰고, 학습이 불안정하면 RAdam 으로 시도해 볼 수 있다.
• 실무에서는 데이터에 따라 ResNet, EfficientNet, ViT 등 다른 백본을 쓰거나, 손실 함수·학습률 스케줄러까지 함께 조정해 나가면 더 좋은 성능을 얻을 수 있다.