train_arcface.py

import pickle
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import warnings
warnings.filterwarnings('ignore')

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("Training device:", device)

# 파일 경로
train_emb_file = 'train_embeddings.pkl'
val_emb_file   = 'val_embeddings.pkl'
label_to_id_file = 'label_to_id.pkl'
model_save_path  = 'classifier.pth'

# 데이터셋 클래스
def l2norm(x): return x / np.linalg.norm(x, axis=1, keepdims=True)

class FaceDataset(Dataset):
    def __init__(self, emb_list, label_list):
        X = np.vstack(emb_list)
        self.X = torch.tensor(l2norm(X), dtype=torch.float32)
        self.y = torch.tensor(label_list, dtype=torch.long)
    def __len__(self): return len(self.y)
    def __getitem__(self, idx): return self.X[idx], self.y[idx]

# 심플 분류기
class SimpleClassifier(nn.Module):
    def __init__(self, embedding_dim=512, num_classes=400):
        super().__init__()
        self.fc = nn.Linear(embedding_dim, num_classes)
    def forward(self, x): return self.fc(x)

# 학습 함수
def train():
    # 로드
    with open(train_emb_file,'rb') as f: train_emb = pickle.load(f)
    with open(val_emb_file,'rb')   as f: val_emb   = pickle.load(f)
    with open(label_to_id_file,'rb') as f: label2id = pickle.load(f)

    X_train, y_train = [], []
    for pid, embs in train_emb.items():
        for e in embs:
            X_train.append(e); y_train.append(label2id[pid])
    X_val, y_val = [], []
    for pid, embs in val_emb.items():
        for e in embs:
            X_val.append(e); y_val.append(label2id[pid])

    train_ds = FaceDataset(X_train, y_train)
    val_ds   = FaceDataset(X_val,   y_val)
    train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
    val_loader   = DataLoader(val_ds, batch_size=64)

    num_classes = len(label2id)
    model = SimpleClassifier(512, num_classes).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    best_val = float('inf')
    for epoch in range(1, 51):
        model.train()
        train_loss = 0
        for X, y in train_loader:
            X, y = X.to(device), y.to(device)
            optimizer.zero_grad()
            loss = criterion(model(X), y)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
        train_loss /= len(train_loader)

        model.eval()
        val_loss = 0
        with torch.no_grad():
            for X, y in val_loader:
                X, y = X.to(device), y.to(device)
                val_loss += criterion(model(X), y).item()
        val_loss /= len(val_loader)

        print(f"Epoch {epoch}: train={train_loss:.4f}, val={val_loss:.4f}")
        if val_loss < best_val:
            best_val = val_loss
            torch.save({'model': model.state_dict(), 'epoch': epoch, 'val_loss': val_loss}, model_save_path)
            print("Model saved")

    print("Training complete!")

if __name__ == '__main__':
    train()