[Pytorch] Pre-traing Vision Transformer로 Fine-tuning : 이미지 분류기

🚀 Model 
- Pytorch에서 제공하는 ViT model 사용(Vit_b_16)
- Link : Pre-trained weights, github code

🚀 Dataset
- Kaggel의 Animal-10 dataset을 강아지와 고양이 이미지만 뽑아서 사용

🚀 Task
- 강아지와 고양이 이미지를 분류

 

1. 전체 실행 흐름 (Main)

전체 코드의 흐름은 아래와 같다.

def main():    
    # 데이터 준비
    FINE_TUNE_N = 1024   # 각 클래스별로 fine tuning에 사용할 이미지 수
    PREDICT_N = 100    # 예측에 사용할 전체 이미지 수
    prepare_fine_tuning_dataset(animal_dir, fine_tuning_dir, FINE_TUNE_N)
    prepare_prediction_dataset(animal_dir, predict_dir, PREDICT_N)
    
    # 모델 fine tuning
    model = fine_tune_model(fine_tuning_dir, num_epochs=20, batch_size=32, learning_rate=1e-4)
    
    # 예측 및 로그 저장
    predict_model(model, predict_dir, batch_size=8)
    
    # 로그 비교 및 평가
    evaluate_predictions()

 

① kaggle의 dataset을 받아와 훈련용 데이터와 테스트용 데이터를 랜덤으로 뽑아 분리

- 기존 model은 Imagenet으로 Pre-train 되어 있기 때문에 본 코드의 task인 강아지와 고양이 이미지를 분류하기 위해 필요함

- 데이터는 basic data에서 랜덤으로 뽑아와 train과 test를 할 수 있도록 코드 구성

- basic data는 cat 폴더와 dog 폴더로 구성이 되어있기 때문에 이미지에 대한 라벨을 적는 log.txt로 학습 및 평가 진행

 

 

② 훈련용 데이터로 Fine tuning 진행

- Main에서 설정한 값을 기준으로 사전 학습된 Vision Transformer를 불러와 Fine-tuning 진행

 

③ 테스트용 데이털 Test 진행 후, 정확도 측정

 

2. 결과 확인

아래의 테스트 train data = 1024, test data = 100으로 설정하여 분류 결과를 확인하였다.

두 번의 측정 모두 분류가 잘되는 것으로 확인되며, Fine tuning이 잘 먹여진 것을 확인할 수 있다.

(초반에 정확도가 너무 낮아서 훈련 데이터를 점차 늘렸음)

 

3. 전체 코드

import os
import glob
import random
import shutil
import re
from PIL import Image

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.models import vit_b_16, ViT_B_16_Weights

# 폴더 없으면 생성하는 함수
def create_dir(path):
    if not os.path.exists(path):
        os.makedirs(path)

# [Step 3] Fine Tuning 데이터셋 준비 (각 클래스에서 N개씩 복사)
def prepare_fine_tuning_dataset(animal_dir, fine_tuning_dir, N):
    categories = ["cat", "dog"]
    for category in categories:
        category_dir = os.path.join(animal_dir, category)
        images = glob.glob(os.path.join(category_dir, "*.jpeg"))
        if len(images) < N:
            print(f"{category_dir}에 이미지가 부족합니다. (필요:{N}, 있음:{len(images)})")
            continue
        selected = random.sample(images, N)
        for i, img_path in enumerate(selected, start=1):
            filename = f"{category}_Img{i}.jpeg"
            dest_path = os.path.join(fine_tuning_dir, filename)
            shutil.copy(img_path, dest_path)
    print("Fine tuning 데이터셋 준비 완료.")

# [Step 4] 예측 데이터셋 준비 및 DataLog.txt 생성
def prepare_prediction_dataset(animal_dir, predict_dir, N):
    categories = ["cat", "dog"]
    all_images = []
    for category in categories:
        category_dir = os.path.join(animal_dir, category)
        images = glob.glob(os.path.join(category_dir, "*.jpeg"))
        for img_path in images:
            all_images.append((img_path, category))
    if len(all_images) < N:
        print(f"전체 이미지가 부족합니다. (필요:{N}, 있음:{len(all_images)})")
        N = len(all_images)
    selected = random.sample(all_images, N)
    random.shuffle(selected)  # 순서를 섞음
    data_log_path = os.path.join("Data", "DataLog.txt")
    with open(data_log_path, "w") as f:
        for i, (img_path, category) in enumerate(selected, start=1):
            filename = f"Img{i}.jpeg"
            dest_path = os.path.join(predict_dir, filename)
            shutil.copy(img_path, dest_path)
            f.write(f"{filename}: {category}\n")
    print("Predict 데이터셋 준비 완료 및 DataLog.txt 생성됨.")

# Fine Tuning용 커스텀 데이터셋 (파일명으로부터 label 추출)
class FineTuneDataset(Dataset):
    def __init__(self, root, transform=None):
        self.root = root
        self.transform = transform
        self.image_files = [f for f in os.listdir(root) if f.endswith(".jpeg")]
        
    def __len__(self):
        return len(self.image_files)
    
    def __getitem__(self, idx):
        filename = self.image_files[idx]
        # 파일명 형식: "cat_Img{i}.jpeg" 또는 "dog_Img{i}.jpeg"
        label_str = filename.split("_")[0]
        label = 0 if label_str.lower() == "cat" else 1
        img_path = os.path.join(self.root, filename)
        image = Image.open(img_path).convert("RGB")
        if self.transform:
            image = self.transform(image)
        return image, label

# 예측용 커스텀 데이터셋 (파일명 정렬)
class PredictDataset(Dataset):
    def __init__(self, root, transform=None):
        self.root = root
        self.transform = transform
        self.image_files = [f for f in os.listdir(root) if f.endswith(".jpeg")]
        self.image_files.sort(key=lambda x: int(re.search(r"(\d+)", x).group(1)))
        
    def __len__(self):
        return len(self.image_files)
    
    def __getitem__(self, idx):
        filename = self.image_files[idx]
        img_path = os.path.join(self.root, filename)
        image = Image.open(img_path).convert("RGB")
        if self.transform:
            image = self.transform(image)
        return image, filename

def fine_tune_model(fine_tuning_dir, num_epochs=10, batch_size=8, learning_rate=1e-4):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
    dataset = FineTuneDataset(fine_tuning_dir, transform=transform)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    
    # ImageNet으로 사전학습된 Vision Transformer 불러오기
    weights = ViT_B_16_Weights.IMAGENET1K_V1
    model = vit_b_16(weights=weights)
    # 분류 head를 2 클래스 분류로 변경 (cat:0, dog:1)
    num_features = model.heads.head.in_features
    model.heads.head = nn.Linear(num_features, 2)
    model = model.to(device)
    
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        for images, labels in dataloader:
            images = images.to(device)
            labels = labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item() * images.size(0)
        
        epoch_loss = running_loss / len(dataset)
        print(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")
        
        # loss가 0.001 이하이면 학습 중단
        if epoch_loss < 0.001:
            print("Loss threshold reached. Stopping training.")
            break

    print("Fine tuning 완료.")
    return model


# [Step 4 ~ 5] 예측 및 PredictLog.txt 기록 함수
def predict_model(model, predict_dir, batch_size=8):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
    dataset = PredictDataset(predict_dir, transform=transform)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
    
    model.eval()
    predictions = {}
    with torch.no_grad():
        for images, filenames in dataloader:
            images = images.to(device)
            outputs = model(images)
            _, preds = torch.max(outputs, 1)
            preds = preds.cpu().numpy()
            for filename, pred in zip(filenames, preds):
                # 숫자 label을 문자열(label)로 변환
                pred_label = "cat" if pred == 0 else "dog"
                predictions[filename] = pred_label
    predict_log_path = os.path.join("Data", "PredictLog.txt")
    with open(predict_log_path, "w") as f:
        # 파일명 순서대로 기록 (Img1.jpeg, Img2.jpeg, …)
        filenames_sorted = sorted(predictions.keys(), key=lambda x: int(re.search(r"(\d+)", x).group(1)))
        for filename in filenames_sorted:
            f.write(f"{filename}: {predictions[filename]}\n")
    print("예측 완료 및 PredictLog.txt 생성됨.")
    
# [Step 6] DataLog.txt와 PredictLog.txt 비교하여 정확도 평가 및 잘못 예측된 이미지 파일 출력
def evaluate_predictions():
    data_log_path = os.path.join("Data", "DataLog.txt")
    predict_log_path = os.path.join("Data", "PredictLog.txt")
    
    with open(data_log_path, "r") as f:
        data_lines = f.readlines()
    with open(predict_log_path, "r") as f:
        predict_lines = f.readlines()
    
    if len(data_lines) != len(predict_lines):
        print("경고: DataLog와 PredictLog의 항목 수가 다릅니다.")
    
    total = min(len(data_lines), len(predict_lines))
    correct = 0
    misclassified = []
    
    for i in range(total):
        # 각 줄은 "ImgX.jpeg: label" 형식
        data_parts = data_lines[i].strip().split(":")
        predict_parts = predict_lines[i].strip().split(":")
        filename = data_parts[0].strip()
        true_label = data_parts[1].strip()
        pred_label = predict_parts[1].strip()
        
        if true_label == pred_label:
            correct += 1
        else:
            misclassified.append((filename, true_label, pred_label))
    
    print(f"정확도: {correct} / {total} (일치하는 개수)")
    
    if misclassified:
        print("\n잘못 예측된 이미지 파일들:")
        for filename, true_label, pred_label in misclassified:
            print(f"{filename}: 실제 = {true_label}, 예측 = {pred_label}")
    else:
        print("\n모든 이미지가 정확하게 예측되었습니다.")

# 전체 실행 흐름
def main():
    # 디렉토리 설정
    base_dir = "Data"
    animal_dir = os.path.join(base_dir, "Animal-10")
    fine_tuning_dir = os.path.join(base_dir, "Fine_tuning")
    predict_dir = os.path.join(base_dir, "Predict")
    
    # 폴더 생성 (없으면)
    create_dir(fine_tuning_dir)
    create_dir(predict_dir)
    
    # 사용할 이미지 개수 (필요에 따라 조정)
    FINE_TUNE_N = 1024   # 각 클래스별로 fine tuning에 사용할 이미지 수
    PREDICT_N = 100    # 예측에 사용할 전체 이미지 수
    
    # 데이터 준비
    prepare_fine_tuning_dataset(animal_dir, fine_tuning_dir, FINE_TUNE_N)
    prepare_prediction_dataset(animal_dir, predict_dir, PREDICT_N)
    
    # 모델 fine tuning
    model = fine_tune_model(fine_tuning_dir, num_epochs=20, batch_size=32, learning_rate=1e-4)
    
    # 예측 및 로그 저장
    predict_model(model, predict_dir, batch_size=8)
    
    # 로그 비교 및 평가
    evaluate_predictions()

if __name__ == '__main__':
    main()