|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
import torch.optim as optim
|
|
|
import os
|
|
|
import argparse
|
|
|
from tqdm import tqdm
|
|
|
from datetime import datetime
|
|
|
from torchvision import transforms
|
|
|
from torch.utils.data import DataLoader
|
|
|
|
|
|
# Import Custom Modules
|
|
|
from model import get_face_model, ArcFace
|
|
|
from dataset import get_face_dataloaders
|
|
|
from validation import LFWDataset, validate_lfw
|
|
|
|
|
|
def main():
|
|
|
# 1. Hyperparameters
|
|
|
parser = argparse.ArgumentParser(description='Face Recognition Training with LFW Validation')
|
|
|
parser.add_argument('--data_dir', type=str, default='/home/cuuva/face_exp/datasets', help='Root dataset dir')
|
|
|
parser.add_argument('--project_dir', type=str, default='./results', help='Save dir')
|
|
|
parser.add_argument('--epochs', type=int, default=20)
|
|
|
parser.add_argument('--batch_size', type=int, default=64)
|
|
|
parser.add_argument('--lr', type=float, default=0.1)
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
print(f"Device: {device}")
|
|
|
|
|
|
# 디렉토리 설정
|
|
|
current_time = datetime.now().strftime("%y-%m-%d-%H-%M")
|
|
|
base_save_dir = os.path.join(args.project_dir, current_time)
|
|
|
best_save_dir = os.path.join(base_save_dir, 'best_model')
|
|
|
os.makedirs(best_save_dir, exist_ok=True)
|
|
|
|
|
|
print(f"✅ Save path: {base_save_dir}")
|
|
|
|
|
|
# 2. Train Data Loader
|
|
|
print("Loading Train Data...")
|
|
|
try:
|
|
|
train_loader, num_classes = get_face_dataloaders(args.data_dir, args.batch_size)
|
|
|
except Exception as e:
|
|
|
print(f"❌ Train Data Error: {e}")
|
|
|
return
|
|
|
|
|
|
# --------------------------------------------------------
|
|
|
# [수정됨] 3. LFW Validation Loader (경로 로직 개선)
|
|
|
# --------------------------------------------------------
|
|
|
print("Loading LFW Data...")
|
|
|
|
|
|
# pairs.txt가 있는 루트 경로: ~/face_exp/datasets/LFW
|
|
|
lfw_root = os.path.join(args.data_dir, 'LFW')
|
|
|
pairs_path = os.path.join(lfw_root, 'pairs.txt')
|
|
|
|
|
|
# 실제 이미지가 있는 경로 찾기
|
|
|
# 1순위: ~/face_exp/datasets/LFW/lfw (보여주신 구조)
|
|
|
# 2순위: ~/face_exp/datasets/LFW (일반적인 구조)
|
|
|
if os.path.exists(os.path.join(lfw_root, 'lfw')):
|
|
|
lfw_img_dir = os.path.join(lfw_root, 'lfw')
|
|
|
else:
|
|
|
lfw_img_dir = lfw_root
|
|
|
|
|
|
print(f"ℹ️ Pairs path: {pairs_path}")
|
|
|
print(f"ℹ️ Images path: {lfw_img_dir}")
|
|
|
|
|
|
if os.path.exists(pairs_path) and os.path.exists(lfw_img_dir):
|
|
|
lfw_transform = transforms.Compose([
|
|
|
transforms.Resize((128, 128), interpolation=transforms.InterpolationMode.BICUBIC),
|
|
|
transforms.ToTensor(),
|
|
|
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
|
|
|
])
|
|
|
|
|
|
# 여기서 lfw_img_dir를 넘겨주는 것이 핵심!
|
|
|
lfw_dataset = LFWDataset(lfw_img_dir, pairs_path, transform=lfw_transform)
|
|
|
|
|
|
lfw_loader = DataLoader(lfw_dataset, batch_size=64, shuffle=False, num_workers=4, drop_last=False)
|
|
|
do_validation = True
|
|
|
print(f"✅ LFW Loaded. Pairs: {len(lfw_dataset)}")
|
|
|
else:
|
|
|
print(f"⚠️ Warning: 'pairs.txt' or Image dir not found. Skipping Validation.")
|
|
|
do_validation = False
|
|
|
|
|
|
# 4. Model & Loss
|
|
|
backbone = get_face_model().to(device)
|
|
|
metric_fc = ArcFace(in_features=128, out_features=num_classes).to(device)
|
|
|
criterion = nn.CrossEntropyLoss()
|
|
|
|
|
|
optimizer = optim.SGD([
|
|
|
{'params': backbone.parameters()},
|
|
|
{'params': metric_fc.parameters()}
|
|
|
], lr=args.lr, momentum=0.9, weight_decay=5e-4)
|
|
|
|
|
|
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[8, 14, 18], gamma=0.1)
|
|
|
|
|
|
# 5. Training Loop
|
|
|
print("🚀 Start Training...")
|
|
|
|
|
|
best_acc = 0.0
|
|
|
|
|
|
for epoch in range(args.epochs):
|
|
|
# --- TRAIN ---
|
|
|
backbone.train()
|
|
|
metric_fc.train()
|
|
|
|
|
|
running_loss = 0.0
|
|
|
pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs}")
|
|
|
|
|
|
for images, labels in pbar:
|
|
|
images, labels = images.to(device), labels.to(device)
|
|
|
optimizer.zero_grad()
|
|
|
|
|
|
features = backbone(images)
|
|
|
outputs = metric_fc(features, labels)
|
|
|
loss = criterion(outputs, labels)
|
|
|
loss.backward()
|
|
|
optimizer.step()
|
|
|
|
|
|
running_loss += loss.item()
|
|
|
pbar.set_postfix({'loss': f"{running_loss / (pbar.n + 1):.4f}"})
|
|
|
|
|
|
scheduler.step()
|
|
|
|
|
|
# --- VALIDATION (LFW) ---
|
|
|
if do_validation:
|
|
|
acc, th = validate_lfw(backbone, lfw_loader, device)
|
|
|
print(f"📊 LFW Acc: {acc*100:.2f}% (Threshold: {th:.2f})")
|
|
|
|
|
|
if acc > best_acc:
|
|
|
best_acc = acc
|
|
|
save_path = os.path.join(best_save_dir, "best_backbone.pt")
|
|
|
torch.save(backbone.state_dict(), save_path)
|
|
|
print(f"🏆 Best Model Updated! Saved to {save_path}")
|
|
|
else:
|
|
|
save_path = os.path.join(base_save_dir, f"backbone_epoch_{epoch+1}.pt")
|
|
|
torch.save(backbone.state_dict(), save_path)
|
|
|
|
|
|
torch.save(backbone.state_dict(), os.path.join(base_save_dir, "last_backbone.pt"))
|
|
|
|
|
|
print("\n🎉 Training Finished!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
main() |
