V1 로직(Linear Probing) - 100-hours-a-week/5-yeosa-wiki GitHub Wiki

1. 로직 설명

  • 이미지의 미적 점수를 판별하기 위해, CLIP 임베딩(ViT-B/32 모델 기준 512차원)을 입력으로 받아 Aesthetic Score를 예측하는 모델 학습
  • 512차원을 입력으로 받아 1차원의 점수를 출력하는 Linear Probing모델 사용

2. 학습 과정

a. 사용 데이터셋

  • AVA Dataset
  • 약 25만장의 이미지 데이터로, 이미지마다 여러명의 사람이 1~10점까지의 점수를 측정한 데이터가 존재

b. 모델 및 설정 조합

항목 구성
CLIP 모델 ViT-B/16, ViT-B/32
score 정규화 여부 True, False
총 조합 수 2 (모델) × 2 (정규화) = 4가지 학습

c. 데이터 전처리

  • score_1~score_10 가중 평균으로 mean_score 산출
  • 정규화 여부에 따라:
    • normed_train / normed_testmean_score / 10.0
    • train / test → 정규화 없음
  • 각 split은 reset_index()로 순서 안정화

d. 이미지 임베딩 저장

  • 모델별 + 정규화 여부별로 .pt 저장 (32개 단위 배치)
  • .pt에는 다음 포함:
    • image_names: 이미지 경로
    • image_features: CLIP 임베딩
    • scores: 정답 aesthetic 점수

e. 학습 설정

항목 내용
모델 구조 AestheticRegressor: nn.Linear(512, 1)
손실 함수 MSELoss
Optimizer Adam(lr=1e-4)
배치 사이즈 32
학습 전략 epoch별 반복, 체크포인트 저장

f. 평가 설정

  • .pt'scores' 그대로 사용 → 안전한 순서 보장
  • 평가 지표: Spearman correlation

e. 사용한 코드

import os
import json
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
from pathlib import Path
from PIL import Image
from scipy.stats import spearmanr
DRIVE_PATH = ''
AVA_PATH = os.path.join(DRIVE_PATH, 'AVA_dataset')
txt_path = os.path.join(AVA_PATH, 'AVA.txt')
dataset_path = os.path.join(AVA_PATH, 'image')

def load_ava_txt(txt_path):
    columns = ['index', 'image_id'] + [f'score_{i}' for i in range(1, 11)] + [f'tag_{i}' for i in range(1, 3)] + ['challenge']
    df = pd.read_csv(txt_path, sep=' ', header=None, names=columns)
    return df
def split_ava_data(data, train_size=229971, norm=True):
    score_bins = np.arange(1, 11)
    score_counts = data[f'score_{i}' for i in range(1, 11)](/100-hours-a-week/5-yeosa-wiki/wiki/f'score_{i}'-for-i-in-range(1,-11)).values
    mean_scores = (score_counts * score_bins).sum(axis=1) / score_counts.sum(axis=1)
    data['image_url'] = data['image_id'].apply(lambda id: os.path.join(dataset_path, f"{id}.jpg"))
    data['mean_score'] = mean_scores if not norm else mean_scores / 10.0

    df = data.sample(frac=1, random_state=42).reset_index(drop=True)
    df_train = df.iloc[:train_size]['image_url', 'mean_score'](/100-hours-a-week/5-yeosa-wiki/wiki/'image_url',-'mean_score')
    df_test = df.iloc[train_size:]['image_url', 'mean_score'](/100-hours-a-week/5-yeosa-wiki/wiki/'image_url',-'mean_score')
    df_test = df_test.reset_index(drop=True)
    df_train = df_train.reset_index(drop=True)
    return df_train, df_test
    
    
normed_train, normed_test = split_ava_data(data)
train, test = split_ava_data(data, norm=False)
import clip

device = 'cuda' if torch.cuda.is_available() else 'cpu'
B_16, B_16_preprocess = clip.load("ViT-B/16", device)
B_32, B_32_preprocess = clip.load("ViT-B/32", device)
def batch_image_generator(df: pd.DataFrame, preprocess, start_batch=0, batch_size: int=32):
    total = len(df)
    for start in range(start_batch*batch_size, total, batch_size):
        count = start // batch_size + 1
        end = min(start + batch_size, total)
        batch_df = df.iloc[start:end]

        images, scores = [], []
        for _, row in batch_df.iterrows():
            try:
                img = Image.open(row['image_url']).convert("RGB")
                img_tensor = preprocess(img)
                images.append(img_tensor)
                scores.append(row['mean_score'])
            except:
                print(f"[Error] Faild to load {row['image_url']}")
                continue
        yield count, torch.stack(images), torch.tensor(scores)
for norm in [True, False]:
	for model_name in ['B_16', 'B_32']:
		train_data = normed_train if norm else train
		model = B_16 if model_name == 'B_16' else B_32
		preprocess = B_16_preprocess if model_name == 'B_16' else B_32_preprocess
		for count, images, scores in batch_image_generator(train_data, preprocess, start_batch=0, batch_size=32):
		    embeddings = embed_images(model, images, device=device)
		
		    start_idx = (count - 1) * 32
		    end_idx = start_idx + len(images)
		    torch.save({
		        'image_names': normed_train.iloc[start_idx:end_idx]['image_url'].tolist(),
		        'image_features': embeddings.cpu(),
		        'scores': scores.cpu(),
		    }, f"ava_embeddings/{'normed_' if norm else ''}ViT_{model_name}_train_image_batch_{count:04d}.pt")
		   
		   
		test_data = normed_test if norm else test
		for count, images, scores in batch_image_generator(test, preprocess, start_batch=0, batch_size=32):
    embeddings = embed_images(model, images, device=device)

    start_idx = (count - 1) * 32
    end_idx = start_idx + len(images)
    torch.save({
        'image_names': test.iloc[start_idx:end_idx]['image_url'].tolist(),
        'image_features': embeddings.cpu(),
        'scores': scores.cpu(),
    }, f"ava_embeddings/{'normed_' if norm else ''}ViT_{model_name}_test_image_batch_{count:04d}.pt")
CHECKPOINT_PATH = os.path.join(AVA_PATH, 'checkpoint')
os.makedirs(CHECKPOINT_PATH, exist_ok=True)

def save_checkpoint(model, optimizer=None, norm=True, model_name='B_16', epoch=1):
    torch.save(model.state_dict(), f"{CHECKPOINT_PATH}/{'normed_' if norm else ''}vit_{model_name}_model_latest_epoch{epoch}")
    if optimizer:
        torch.save(optimizer.state_dict(), f"{CHECKPOINT_PATH}/{'normed_' if norm else ''}vit_{model_name}_optimizer_latest_epoch{epoch}.pth")
# 모델 학습 및 평가 루프: 
# 정규화 여부(norm), 모델 종류(B_16/B_32), epoch(1~5)에 따라 반복 실행
for norm in [True, False]:
    for model_name in ['B_16', 'B_32']:
        for epoch in range(1, 6):

            # 1. ⬇ 학습 데이터 불러오기 (임베딩된 .pt 파일 기준)
            train_dir = Path("ava_embeddings")
            prefix = f"{'normed_' if norm else ''}ViT_{model_name}_train_image_batch_"
            train_files = sorted(train_dir.glob(f"{prefix}*.pt"))

            train_features = []
            train_labels = []

            # 각 배치별 .pt 파일에서 임베딩과 score를 누적
            for file in train_files:
                data = torch.load(file)
                feats = data["image_features"]       # [batch_size, 512]
                scores = data["scores"].float()      # [batch_size]

                train_features.append(feats)
                train_labels.append(scores)

            # 전체 학습 데이터를 하나의 텐서로 결합
            X_train = torch.cat(train_features)
            y_train = torch.cat(train_labels)

            # 2. ⬇ 모델 정의 및 초기화
            regressor = AestheticRegressor(in_dim=512).to(device)
            optimizer = optim.Adam(regressor.parameters(), lr=1e-4)
            criterion = nn.MSELoss()

            # 이전 epoch에서 저장한 모델/옵티마이저 상태가 있다면 이어받기
            if epoch > 1:
                regressor.load_state_dict(torch.load(
                    f"{CHECKPOINT_PATH}/{'normed_' if norm else ''}vit_{model_name}_model_latest_epoch{epoch-1}_continued_optimizer.pth"))
                optimizer.load_state_dict(torch.load(
                    f"{CHECKPOINT_PATH}/{'normed_' if norm else ''}vit_{model_name}_optimizer_latest_epoch{epoch-1}_continued_optimizer.pth"))

            regressor.train()
            batch_size = 32

            # 3. ⬇ 학습 루프 (미니배치 학습)
            for i in tqdm(range(0, len(X_train), batch_size)):
                batch_x = X_train[i:i+batch_size].to(device)
                batch_y = y_train[i:i+batch_size].to(device)

                pred_y = regressor(batch_x)
                loss = criterion(pred_y.squeeze(), batch_y)  # 손실 계산

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                step = i // batch_size + 1
                print(f"{'normed_' if norm else ''}{model_name}[Batch {step}] Loss: {loss.item():.4f}")

                # 매 배치마다 체크포인트 저장 (모델 + 옵티마이저 상태)
                save_checkpoint(regressor, step, loss.item(), optimizer, norm, model_name, epoch)

            # 4. ⬇ 테스트셋 불러오기 (임베딩된 .pt)
            test_dir = Path("ava_embeddings")
            prefix = f"{'normed_' if norm else ''}ViT_{model_name}_test_image_batch_"
            test_files = sorted(test_dir.glob(f"{prefix}*.pt"))

            regressor.eval()
            y_true = []
            y_pred = []

            # 평가 루프
            with torch.no_grad():
                for file in tqdm(test_files):
                    data = torch.load(file)
                    features = data["image_features"].to(device)
                    scores = data["scores"]  # 저장된 ground-truth mean_score

                    pred_scores = regressor(features).squeeze().cpu()

                    y_true.extend(scores.tolist())
                    y_pred.extend(pred_scores.tolist())

            # 5. ⬇ Spearman correlation 평가
            spearman_corr = spearmanr(y_true, y_pred).correlation
            print(f"✅ Spearman correlation: {spearman_corr:.4f}")

            # 결과 저장 디렉토리 생성 (없으면 생성됨)
            result_path = Path("results")
            result_path.mkdir(exist_ok=True)

            # 결과 저장 (.txt 파일) - 설정별로 구분되게 저장됨
            with open(result_path / f"{'normed_' if norm else ''}ViT_{model_name}_epoch_{epoch}_continued_optimizer_spearman_result.txt", "w") as f:
                f.write(f"Spearman correlation: {spearman_corr:.4f}\n")


3. 실험 결과

a. 정규화 유무의 영향

모델 Raw 최고 Spearman Normed 최고 Spearman
B_16 0.5678 0.6844
B_32 0.5250 0.6753
  • 정규화(normed) 처리가 성능을 크게 향상시킴

    → 모델이 작은 수치 범위 내에서 더 안정적으로 학습 가능

b. 모델별 성능 비교

조건 최고 Spearman
ViT-B/16 + Normed 0.6844
ViT-B/32 + Normed 0.6753
ViT-B/16 + Raw 0.5678
ViT-B/32 + Raw 0.5250
  • ViT-B/16이 전체적으로 더 높은 성능을 보임

c. Epoch 별 성능 변화

가. ViT-B/16 성능 변화

Epoch Raw Normed
1 0.2130 0.5783
2 0.3180 0.6551
3 0.4541 0.6735
4 0.5269 0.6808
5 0.5678 ⭐ 0.6844
  • 해석:
    • Normed의 경우 epoch 2부터 급상승 → 이후 완만한 상승
    • Raw는 점진적 상승이지만 최종 성능은 낮음
    • ✅ B/16은 정규화 + 충분한 epoch에서 최고의 성능 발휘

나. ViT-B/32 성능 변화

Epoch Raw Normed
1 0.2103 0.5425
2 0.2766 0.6432
3 0.3940 0.6646
4 0.4748 0.6721
5 0.5250 ⭐ 0.6753
  • 해석:
    • ViT-B/32normed 설정에서 급격한 상승을 보임
    • raw는 비교적 천천히 상승하지만 최종 성능은 낮음
    • 정규화만 하면 B/16과 큰 차이 없이 수렴 가능

다. Epoch별 비교 총평

모델 정규화 Epoch별 향상 추이 특징 요약
B/16 느리지만 꾸준히 상승 최종 성능 낮음
B/16 Epoch 2~3에서 급격한 상승 빠른 수렴 + 최고 성능
B/32 느림, 최종 0.525 불리함
B/32 Epoch 2~3에서 급상승 안정적 수렴, 최고 0.675

실험 결과 비교2.png