Spaces:

Harshasnade
/

Deepfake_Detection_System_V1

Running

Deepfake_Detection_System_V1 / model /src /train.py

Deploy Backend (No Frontend)

0966609 10 months ago

11.5 kB

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm
import random
import ssl
# Disable SSL verification for downloading pretrained weights
ssl._create_default_https_context = ssl._create_unverified_context
from torch.cuda.amp import GradScaler, autocast

from src.config import Config
from src.models import DeepfakeDetector
from src.dataset import DeepfakeDataset

try:
    from safetensors.torch import save_file, load_file
    SAFETENSORS_AVAILABLE = True
except ImportError:
    SAFETENSORS_AVAILABLE = False
    print("Warning: safetensors not installed. Checkpoints will be saved as .pt")

def train():
    # Setup
    Config.setup()
    device = torch.device(Config.DEVICE)
    
    # --- Data Loading with Automatic Split ---
    if Config.TRAIN_DATA_PATH == Config.TEST_DATA_PATH:
        print("Train and Test paths are identical. Performing automatic 80/20 shuffle split...")
        all_paths, all_labels = DeepfakeDataset.scan_directory(Config.TRAIN_DATA_PATH)
        
        if len(all_paths) == 0:
            print(f"No images found in {Config.TRAIN_DATA_PATH}")
            return

        # Combine and shuffle
        combined = list(zip(all_paths, all_labels))
        random.shuffle(combined)
        
        split_idx = int(len(combined) * 0.8)
        train_data = combined[:split_idx]
        val_data = combined[split_idx:]
        
        train_paths, train_labels = zip(*train_data)
        val_paths, val_labels = zip(*val_data)
        
        train_dataset = DeepfakeDataset(file_paths=list(train_paths), labels=list(train_labels), phase='train')
        val_dataset = DeepfakeDataset(file_paths=list(val_paths), labels=list(val_labels), phase='val')
    else:
        # Standard folder-based loading
        train_dataset = DeepfakeDataset(root_dir=Config.TRAIN_DATA_PATH, phase='train')
        val_dataset = DeepfakeDataset(root_dir=Config.TEST_DATA_PATH, phase='val')
    
    # Dataloaders
    train_loader = DataLoader(train_dataset, batch_size=Config.BATCH_SIZE, shuffle=True,
                              num_workers=Config.NUM_WORKERS, 
                              pin_memory=True if device.type=='cuda' else False,
                              persistent_workers=True if Config.NUM_WORKERS > 0 else False)
    val_loader = DataLoader(val_dataset, batch_size=Config.BATCH_SIZE, shuffle=False,
                            num_workers=Config.NUM_WORKERS, 
                            pin_memory=True if device.type=='cuda' else False,
                            persistent_workers=True if Config.NUM_WORKERS > 0 else False)
    
    # Model
    print("Initializing Multi-Branch DeepfakeDetector...")
    model = DeepfakeDetector(pretrained=True).to(device)
    
    # Optimization
    criterion = nn.BCEWithLogitsLoss()
    optimizer = optim.AdamW(model.parameters(), lr=Config.LEARNING_RATE, weight_decay=Config.WEIGHT_DECAY)
    # Optimization
    criterion = nn.BCEWithLogitsLoss()
    optimizer = optim.AdamW(model.parameters(), lr=Config.LEARNING_RATE, weight_decay=Config.WEIGHT_DECAY)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
    
    # Enable AMP only for CUDA (Windows NVIDIA)
    use_amp = (Config.DEVICE == 'cuda')
    scaler = GradScaler() if use_amp else None
    if use_amp:
        print("🚀 Mixed Precision (AMP) Enabled for RTX GPU")
    else:
        print("🐌 Standard Precision (No AMP) for CPU/MPS")
    
    # Resume from checkpoint if exists
    start_epoch = 0
    best_acc = 0.0
    
    # Priority:
    # 1. best_model.safetensors (if we crashed mid-training)
    # 2. patched_model.safetensors (the model we want to improve)
    
    resume_path = os.path.join(Config.CHECKPOINT_DIR, "best_model.safetensors")
    if not os.path.exists(resume_path):
        # Look for latest epoch checkpoint
        import glob
        import re
        checkpoints = glob.glob(os.path.join(Config.CHECKPOINT_DIR, "checkpoint_ep*.safetensors"))
        if checkpoints:
            # Sort by epoch number
            def get_epoch(p):
                match = re.search(r"checkpoint_ep(\d+)", p)
                return int(match.group(1)) if match else 0
            
            latest_ckpt = max(checkpoints, key=get_epoch)
            resume_path = latest_ckpt
            start_epoch = get_epoch(latest_ckpt)
            print(f"🔄 Auto-Resuming from latest epoch: {start_epoch}")
        else:
            resume_path = os.path.join(Config.CHECKPOINT_DIR, "patched_model.safetensors")
    
    if os.path.exists(resume_path):
        print(f"\n🔄 Found existing checkpoint: {resume_path}")
        print("Auto-resuming to FINETUNE this model...")
        
        try:
            if resume_path.endswith(".safetensors") and SAFETENSORS_AVAILABLE:
                state_dict = load_file(resume_path)
            else:
                state_dict = torch.load(resume_path, map_location=device)
            
            # Use strict=False to allow for minor architecture changes or missing keys
            model.load_state_dict(state_dict, strict=False)
            print("✅ Weights loaded. Starting Fine-Tuning.")
        except Exception as e:
            print(f"⚠ Failed to load checkpoint: {e}")
            print("Starting from ImageNet weights.")
    
    # Loop
    
    for epoch in range(start_epoch, Config.EPOCHS):
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0
        
        loop = tqdm(train_loader, desc=f"Epoch {epoch+1}/{Config.EPOCHS}")
        for images, labels in loop:
            images = images.to(device)
            labels = labels.to(device).unsqueeze(1)
            
            optimizer.zero_grad()
            
            if use_amp:
                with autocast():
                    outputs = model(images)
                    loss = criterion(outputs, labels)
                
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
            else:
                # Standard training for Mac/CPU
                outputs = model(images)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()
            
            train_loss += loss.item()
            preds = (torch.sigmoid(outputs) > 0.5).float()
            correct = (preds == labels).sum().item()
            train_correct += correct
            train_total += labels.size(0)
            
            loop.set_postfix(loss=loss.item(), acc=correct/labels.size(0))
            
        train_acc = train_correct / train_total if train_total > 0 else 0
        print(f"Epoch {epoch+1} Train Loss: {train_loss/len(train_loader):.4f} Acc: {train_acc:.4f}")
        
        # Save checkpoint after every epoch
        save_checkpoint(model, epoch+1, train_acc, best=False)
        
        # Validation
        if len(val_dataset) > 0:
            val_loss, val_acc = validate(model, val_loader, criterion, device)
            print(f"Epoch {epoch+1} Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}")
            
            # Save best model if validation accuracy improved
            if val_acc > best_acc:
                best_acc = val_acc
                print(f"⭐ New best model! Validation Accuracy: {val_acc:.4f}")
                save_checkpoint(model, epoch+1, val_acc, best=True)
        
        scheduler.step()
    
    print(f"\n🎉 Training Complete!")
    print(f"Best Validation Accuracy: {best_acc:.4f}")

def validate(model, loader, criterion, device):
    model.eval()
    val_loss = 0.0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for images, labels in loader:
            images = images.to(device)
            labels = labels.to(device).unsqueeze(1)
            
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            val_loss += loss.item()
            preds = (torch.sigmoid(outputs) > 0.5).float()
            correct += (preds == labels).sum().item()
            total += labels.size(0)
            
    return val_loss / len(loader), correct / total

def save_checkpoint(model, epoch, acc, best=False):
    state_dict = model.state_dict()
    name = "best_model.safetensors" if best else f"checkpoint_ep{epoch}.safetensors"
    path = os.path.join(Config.CHECKPOINT_DIR, name)
    
    if SAFETENSORS_AVAILABLE:
        try:
            # Try with shared tensors support
            from safetensors.torch import save_model
            save_model(model, path)
            print(f"Saved Checkpoint: {path}")
            
            # 📝 Auto-Log to History
            try:
                from datetime import datetime
                log_path = os.path.join(Config.PROJECT_ROOT, "TRAINING_HISTORY.md")
                timestamp = datetime.now().strftime("%Y-%m-%d | %I:%M %p")
                
                # Create file with header if doesn't exist
                if not os.path.exists(log_path):
                    with open(log_path, "w", encoding="utf-8") as f:
                        f.write("# 📜 Training History Log\n\n")
                        f.write("| Date | Time | Model Name | Dataset | Epochs | Accuracy | Loss | Status |\n")
                        f.write("| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n")
                
                # Append Entry to Summary Log
                with open(log_path, "a", encoding="utf-8") as f:
                    # Format: Date | Time | Name | Dataset | Epoch | Acc | Loss | Status
                    dataset_name = os.path.basename(Config.DATA_DIR)
                    entry = f"| **{timestamp.split(' | ')[0]}** | {timestamp.split(' | ')[1]} | {name} | {dataset_name} | {epoch} | {acc*100:.2f}% | N/A | ✅ Saved |\n"
                    f.write(entry)
                    print(f"📝 Logged to TRAINING_HISTORY.md")

                # 📝 Detailed Lab Notebook Logging
                detail_path = os.path.join(Config.PROJECT_ROOT, "DETAILED_HISTORY.md")
                with open(detail_path, "a", encoding="utf-8") as f:
                    f.write(f"\n## Model: {name} (Epoch {epoch})\n")
                    f.write(f"| Feature | Detail |\n| :--- | :--- |\n")
                    f.write(f"| **Date** | {timestamp} |\n")
                    f.write(f"| **Training Accuracy** | {acc*100:.2f}% |\n")
                    f.write(f"| **Dataset** | {Config.DATA_DIR} |\n")
                    f.write(f"| **Batch Size** | {Config.BATCH_SIZE} |\n")
                    f.write(f"| **Optimizer** | AdamW (lr={Config.LEARNING_RATE}) |\n")
                    f.write(f"| **Device** | {Config.DEVICE.upper()} |\n")
                    f.write("\n---\n")
                    print(f"📘 Detailed log written to DETAILED_HISTORY.md")

            except Exception as e:
                print(f"⚠️ Failed to write log: {e}")

        except Exception as e:
            # Fallback to regular torch save if safetensors fails
            print(f"SafeTensors save failed ({e}), falling back to .pth format")
            torch.save(state_dict, path.replace(".safetensors", ".pth"))
            print(f"Saved Checkpoint (Legacy): {path.replace('.safetensors', '.pth')}")
    else:
        torch.save(state_dict, path.replace(".safetensors", ".pth"))
        print(f"Saved Checkpoint (Legacy): {path}")

if __name__ == "__main__":
    train()

	import os
	import torch
	import torch.nn as nn
	import torch.optim as optim
	from torch.utils.data import DataLoader
	from tqdm import tqdm
	import random
	import ssl
	# Disable SSL verification for downloading pretrained weights
	ssl._create_default_https_context = ssl._create_unverified_context
	from torch.cuda.amp import GradScaler, autocast

	from src.config import Config
	from src.models import DeepfakeDetector
	from src.dataset import DeepfakeDataset

	try:
	from safetensors.torch import save_file, load_file
	SAFETENSORS_AVAILABLE = True
	except ImportError:
	SAFETENSORS_AVAILABLE = False
	print("Warning: safetensors not installed. Checkpoints will be saved as .pt")

	def train():
	# Setup
	Config.setup()
	device = torch.device(Config.DEVICE)

	# --- Data Loading with Automatic Split ---
	if Config.TRAIN_DATA_PATH == Config.TEST_DATA_PATH:
	print("Train and Test paths are identical. Performing automatic 80/20 shuffle split...")
	all_paths, all_labels = DeepfakeDataset.scan_directory(Config.TRAIN_DATA_PATH)

	if len(all_paths) == 0:
	print(f"No images found in {Config.TRAIN_DATA_PATH}")
	return

	# Combine and shuffle
	combined = list(zip(all_paths, all_labels))
	random.shuffle(combined)

	split_idx = int(len(combined) * 0.8)
	train_data = combined[:split_idx]
	val_data = combined[split_idx:]

	train_paths, train_labels = zip(*train_data)
	val_paths, val_labels = zip(*val_data)

	train_dataset = DeepfakeDataset(file_paths=list(train_paths), labels=list(train_labels), phase='train')
	val_dataset = DeepfakeDataset(file_paths=list(val_paths), labels=list(val_labels), phase='val')
	else:
	# Standard folder-based loading
	train_dataset = DeepfakeDataset(root_dir=Config.TRAIN_DATA_PATH, phase='train')
	val_dataset = DeepfakeDataset(root_dir=Config.TEST_DATA_PATH, phase='val')

	# Dataloaders
	train_loader = DataLoader(train_dataset, batch_size=Config.BATCH_SIZE, shuffle=True,
	num_workers=Config.NUM_WORKERS,
	pin_memory=True if device.type=='cuda' else False,
	persistent_workers=True if Config.NUM_WORKERS > 0 else False)
	val_loader = DataLoader(val_dataset, batch_size=Config.BATCH_SIZE, shuffle=False,
	num_workers=Config.NUM_WORKERS,
	pin_memory=True if device.type=='cuda' else False,
	persistent_workers=True if Config.NUM_WORKERS > 0 else False)

	# Model
	print("Initializing Multi-Branch DeepfakeDetector...")
	model = DeepfakeDetector(pretrained=True).to(device)

	# Optimization
	criterion = nn.BCEWithLogitsLoss()
	optimizer = optim.AdamW(model.parameters(), lr=Config.LEARNING_RATE, weight_decay=Config.WEIGHT_DECAY)
	# Optimization
	criterion = nn.BCEWithLogitsLoss()
	optimizer = optim.AdamW(model.parameters(), lr=Config.LEARNING_RATE, weight_decay=Config.WEIGHT_DECAY)
	scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

	# Enable AMP only for CUDA (Windows NVIDIA)
	use_amp = (Config.DEVICE == 'cuda')
	scaler = GradScaler() if use_amp else None
	if use_amp:
	print("🚀 Mixed Precision (AMP) Enabled for RTX GPU")
	else:
	print("🐌 Standard Precision (No AMP) for CPU/MPS")

	# Resume from checkpoint if exists
	start_epoch = 0
	best_acc = 0.0

	# Priority:
	# 1. best_model.safetensors (if we crashed mid-training)
	# 2. patched_model.safetensors (the model we want to improve)

	resume_path = os.path.join(Config.CHECKPOINT_DIR, "best_model.safetensors")
	if not os.path.exists(resume_path):
	# Look for latest epoch checkpoint
	import glob
	import re
	checkpoints = glob.glob(os.path.join(Config.CHECKPOINT_DIR, "checkpoint_ep*.safetensors"))
	if checkpoints:
	# Sort by epoch number
	def get_epoch(p):
	match = re.search(r"checkpoint_ep(\d+)", p)
	return int(match.group(1)) if match else 0

	latest_ckpt = max(checkpoints, key=get_epoch)
	resume_path = latest_ckpt
	start_epoch = get_epoch(latest_ckpt)
	print(f"🔄 Auto-Resuming from latest epoch: {start_epoch}")
	else:
	resume_path = os.path.join(Config.CHECKPOINT_DIR, "patched_model.safetensors")

	if os.path.exists(resume_path):
	print(f"\n🔄 Found existing checkpoint: {resume_path}")
	print("Auto-resuming to FINETUNE this model...")

	try:
	if resume_path.endswith(".safetensors") and SAFETENSORS_AVAILABLE:
	state_dict = load_file(resume_path)
	else:
	state_dict = torch.load(resume_path, map_location=device)

	# Use strict=False to allow for minor architecture changes or missing keys
	model.load_state_dict(state_dict, strict=False)
	print("✅ Weights loaded. Starting Fine-Tuning.")
	except Exception as e:
	print(f"⚠ Failed to load checkpoint: {e}")
	print("Starting from ImageNet weights.")

	# Loop

	for epoch in range(start_epoch, Config.EPOCHS):
	model.train()
	train_loss = 0.0
	train_correct = 0
	train_total = 0

	loop = tqdm(train_loader, desc=f"Epoch {epoch+1}/{Config.EPOCHS}")
	for images, labels in loop:
	images = images.to(device)
	labels = labels.to(device).unsqueeze(1)

	optimizer.zero_grad()

	if use_amp:
	with autocast():
	outputs = model(images)
	loss = criterion(outputs, labels)

	scaler.scale(loss).backward()
	scaler.step(optimizer)
	scaler.update()
	else:
	# Standard training for Mac/CPU
	outputs = model(images)
	loss = criterion(outputs, labels)
	loss.backward()
	optimizer.step()

	train_loss += loss.item()
	preds = (torch.sigmoid(outputs) > 0.5).float()
	correct = (preds == labels).sum().item()
	train_correct += correct
	train_total += labels.size(0)

	loop.set_postfix(loss=loss.item(), acc=correct/labels.size(0))

	train_acc = train_correct / train_total if train_total > 0 else 0
	print(f"Epoch {epoch+1} Train Loss: {train_loss/len(train_loader):.4f} Acc: {train_acc:.4f}")

	# Save checkpoint after every epoch
	save_checkpoint(model, epoch+1, train_acc, best=False)

	# Validation
	if len(val_dataset) > 0:
	val_loss, val_acc = validate(model, val_loader, criterion, device)
	print(f"Epoch {epoch+1} Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}")

	# Save best model if validation accuracy improved
	if val_acc > best_acc:
	best_acc = val_acc
	print(f"⭐ New best model! Validation Accuracy: {val_acc:.4f}")
	save_checkpoint(model, epoch+1, val_acc, best=True)

	scheduler.step()

	print(f"\n🎉 Training Complete!")
	print(f"Best Validation Accuracy: {best_acc:.4f}")

	def validate(model, loader, criterion, device):
	model.eval()
	val_loss = 0.0
	correct = 0
	total = 0

	with torch.no_grad():
	for images, labels in loader:
	images = images.to(device)
	labels = labels.to(device).unsqueeze(1)

	outputs = model(images)
	loss = criterion(outputs, labels)

	val_loss += loss.item()
	preds = (torch.sigmoid(outputs) > 0.5).float()
	correct += (preds == labels).sum().item()
	total += labels.size(0)

	return val_loss / len(loader), correct / total

	def save_checkpoint(model, epoch, acc, best=False):
	state_dict = model.state_dict()
	name = "best_model.safetensors" if best else f"checkpoint_ep{epoch}.safetensors"
	path = os.path.join(Config.CHECKPOINT_DIR, name)

	if SAFETENSORS_AVAILABLE:
	try:
	# Try with shared tensors support
	from safetensors.torch import save_model
	save_model(model, path)
	print(f"Saved Checkpoint: {path}")

	# 📝 Auto-Log to History
	try:
	from datetime import datetime
	log_path = os.path.join(Config.PROJECT_ROOT, "TRAINING_HISTORY.md")
	timestamp = datetime.now().strftime("%Y-%m-%d \| %I:%M %p")

	# Create file with header if doesn't exist
	if not os.path.exists(log_path):
	with open(log_path, "w", encoding="utf-8") as f:
	f.write("# 📜 Training History Log\n\n")
	f.write("\| Date \| Time \| Model Name \| Dataset \| Epochs \| Accuracy \| Loss \| Status \|\n")
	f.write("\| :--- \| :--- \| :--- \| :--- \| :--- \| :--- \| :--- \| :--- \|\n")

	# Append Entry to Summary Log
	with open(log_path, "a", encoding="utf-8") as f:
	# Format: Date \| Time \| Name \| Dataset \| Epoch \| Acc \| Loss \| Status
	dataset_name = os.path.basename(Config.DATA_DIR)
	entry = f"\| {timestamp.split(' \| ')[0]} \| {timestamp.split(' \| ')[1]} \| {name} \| {dataset_name} \| {epoch} \| {acc*100:.2f}% \| N/A \| ✅ Saved \|\n"
	f.write(entry)
	print(f"📝 Logged to TRAINING_HISTORY.md")

	# 📝 Detailed Lab Notebook Logging
	detail_path = os.path.join(Config.PROJECT_ROOT, "DETAILED_HISTORY.md")
	with open(detail_path, "a", encoding="utf-8") as f:
	f.write(f"\n## Model: {name} (Epoch {epoch})\n")
	f.write(f"\| Feature \| Detail \|\n\| :--- \| :--- \|\n")
	f.write(f"\| Date \| {timestamp} \|\n")
	f.write(f"\| Training Accuracy \| {acc*100:.2f}% \|\n")
	f.write(f"\| Dataset \| {Config.DATA_DIR} \|\n")
	f.write(f"\| Batch Size \| {Config.BATCH_SIZE} \|\n")
	f.write(f"\| Optimizer \| AdamW (lr={Config.LEARNING_RATE}) \|\n")
	f.write(f"\| Device \| {Config.DEVICE.upper()} \|\n")
	f.write("\n---\n")
	print(f"📘 Detailed log written to DETAILED_HISTORY.md")

	except Exception as e:
	print(f"⚠️ Failed to write log: {e}")

	except Exception as e:
	# Fallback to regular torch save if safetensors fails
	print(f"SafeTensors save failed ({e}), falling back to .pth format")
	torch.save(state_dict, path.replace(".safetensors", ".pth"))
	print(f"Saved Checkpoint (Legacy): {path.replace('.safetensors', '.pth')}")
	else:
	torch.save(state_dict, path.replace(".safetensors", ".pth"))
	print(f"Saved Checkpoint (Legacy): {path}")

	if __name__ == "__main__":
	train()