training/train_efficientnet_auth.py

import sys
import os

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if ROOT_DIR not in sys.path:
    sys.path.insert(0, ROOT_DIR)

# Redirect HF cache to D: drive (C: has no space)
os.environ["HF_HOME"] = os.path.join(ROOT_DIR, ".hf_cache")
os.environ["HF_DATASETS_CACHE"] = os.path.join(ROOT_DIR, ".hf_cache", "datasets")

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import CosineAnnealingLR
from torchvision import transforms

from core_models.efficientnet_auth_model import EfficientNetAuthModel
from training.dataset_portraits import load_portrait_dataset, PortraitDataset

# ---------------- CONFIG ----------------
BATCH_SIZE = 32
EPOCHS = 15
BACKBONE_LR = 1e-5
HEAD_LR = 1e-3
TRAIN_SPLIT = 0.85
MAX_SAMPLES = 20000          # 10x increase for meaningful training
MODEL_PATH = "models/efficientnet_auth_model.pth"
EARLY_STOPPING_PATIENCE = 5
LABEL_SMOOTHING = 0.05
# ---------------------------------------


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Using device:", device)

    # -------- Transforms (video-like augmentation) --------
    train_transform = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(10),
        transforms.RandomApply([transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 2.0))], p=0.3),
        transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.05),
        transforms.RandomGrayscale(p=0.05),
        transforms.ToTensor(),
        transforms.RandomErasing(p=0.1, scale=(0.02, 0.1)),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])

    val_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]
        )
    ])

    # -------- Load dataset via the shared multi-source loader --------
    # Now pulls from PORTRAIT_SOURCES (many GAN/diffusion HF-streaming
    # sources, reservoir-sampled) PLUS OpenRL/DeepFakeFace (diffusion-
    # generated: SD Inpainting/text2img, InsightFace) instead of the single
    # Hemg/AI-Generated-vs-Real-Images-Datasets source used before - that
    # source alone taught this model nothing about diffusion output.
    # training/eval_image_benchmark.py found the whole ensemble at 30.8%
    # accuracy on diffusion fakes vs 83-97% on real photos, and this
    # model's own average score on that held-out set was just 0.153,
    # meaning it was still confidently calling diffusion fakes "real."
    # face_align=False: tried True on the theory that InsightFace fakes
    # (face-swaps onto the same real IMDB-WIKI photos as the real baseline -
    # see training/diagnose_insight.py) would be easier to catch with a
    # tight face crop isolating the manipulated region. Measured result was
    # the opposite: InsightFace accuracy dropped 6.7%->3.3% and inpainting
    # dropped too (training/eval_image_benchmark.py, round 4 vs round 3) -
    # a tight crop likely excludes the blend boundary itself, which sits at
    # the hairline/jaw/neck edge just outside it, and losing the
    # complementary whole-image signal across 3 of 7 fusion inputs hurt
    # more than the crop helped. Reverted.
    print(f"Loading multi-source portrait dataset ({MAX_SAMPLES} samples)...")
    train_data, val_data = load_portrait_dataset(
        max_samples=MAX_SAMPLES, train_split=TRAIN_SPLIT, face_align=False,
    )

    print(f"Train samples: {len(train_data)}")
    print(f"Val samples  : {len(val_data)}")

    train_loader = DataLoader(
        PortraitDataset(train_data, train_transform),
        batch_size=BATCH_SIZE, shuffle=True, num_workers=0, pin_memory=True,
    )
    val_loader = DataLoader(
        PortraitDataset(val_data, val_transform),
        batch_size=BATCH_SIZE, shuffle=False, num_workers=0, pin_memory=True,
    )

    # -------- Model --------
    model = EfficientNetAuthModel().to(device)
    criterion = nn.BCELoss()

    # Differential learning rates
    backbone_params = [p for p in model.model.features.parameters() if p.requires_grad]
    head_params = list(model.model.classifier.parameters())

    optimizer = torch.optim.AdamW([
        {"params": backbone_params, "lr": BACKBONE_LR, "weight_decay": 1e-4},
        {"params": head_params, "lr": HEAD_LR, "weight_decay": 1e-4},
    ])

    scheduler = CosineAnnealingLR(optimizer, T_max=EPOCHS, eta_min=1e-7)

    # -------- Training --------
    print(f"\nStarting EfficientNetV2 training (label_smoothing={LABEL_SMOOTHING})...\n")

    best_val_loss = float("inf")
    patience_counter = 0

    for epoch in range(EPOCHS):
        model.train()
        total_loss = 0.0

        for imgs, labels in train_loader:
            imgs = imgs.to(device)
            labels = labels.unsqueeze(1).to(device)
            smoothed = labels * (1 - LABEL_SMOOTHING) + (1 - labels) * LABEL_SMOOTHING

            preds = model(imgs)
            loss = criterion(preds, smoothed)

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

            total_loss += loss.item()

        avg_loss = total_loss / len(train_loader)

        # -------- Validation --------
        model.eval()
        val_loss = 0.0
        correct, total = 0, 0

        with torch.no_grad():
            for imgs, labels in val_loader:
                imgs = imgs.to(device)
                labels_t = labels.unsqueeze(1).to(device)

                preds = model(imgs)
                loss = criterion(preds, labels_t)
                val_loss += loss.item()

                pred_labels = (preds > 0.5).float()
                correct += (pred_labels == labels_t).sum().item()
                total += labels.size(0)

        avg_val_loss = val_loss / len(val_loader)
        val_acc = correct / total

        scheduler.step()

        print(
            f"Epoch [{epoch+1}/{EPOCHS}] "
            f"| Train Loss: {avg_loss:.4f} "
            f"| Val Loss: {avg_val_loss:.4f} "
            f"| Val Acc: {val_acc:.4f}"
        )

        # -------- Early stopping --------
        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            patience_counter = 0
            os.makedirs("models", exist_ok=True)
            torch.save(model.state_dict(), MODEL_PATH)
            print(f"  -> Best model saved (val_loss={avg_val_loss:.4f})")
        else:
            patience_counter += 1
            if patience_counter >= EARLY_STOPPING_PATIENCE:
                print(f"\nEarly stopping at epoch {epoch+1} (no improvement for {EARLY_STOPPING_PATIENCE} epochs)")
                break

    print("\nEfficientNetV2 authenticity model trained successfully.")
    print(f"Best model saved to: {MODEL_PATH}")


if __name__ == "__main__":
    main()