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_test
→ mean_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 |
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/32
도 normed 설정에서 급격한 상승을 보임
- raw는 비교적 천천히 상승하지만 최종 성능은 낮음
- 정규화만 하면
B/16
과 큰 차이 없이 수렴 가능
다. Epoch별 비교 총평
모델 |
정규화 |
Epoch별 향상 추이 |
특징 요약 |
B/16 |
❌ |
느리지만 꾸준히 상승 |
최종 성능 낮음 |
B/16 |
✅ |
Epoch 2~3에서 급격한 상승 |
빠른 수렴 + 최고 성능 |
B/32 |
❌ |
느림, 최종 0.525 |
불리함 |
B/32 |
✅ |
Epoch 2~3에서 급상승 |
안정적 수렴, 최고 0.675 |
%201e554217ba3f80b69752f4be51a65292/%E1%84%89%E1%85%B5%E1%86%AF%E1%84%92%E1%85%A5%E1%86%B7_%E1%84%80%E1%85%A7%E1%86%AF%E1%84%80%E1%85%AA_%E1%84%87%E1%85%B5%E1%84%80%E1%85%AD2.png)