Delete train.py

train.py

@@ -1,100 +0,0 @@
-# train.py
-import os
-from torchvision import datasets, transforms, models
-import torch
-from torch import nn, optim
-from torch.utils.data import DataLoader
-
-# -----------------------------
-# CONFIG
-# -----------------------------
-DATA_DIR = "datasets"
-BATCH_SIZE = 16
-NUM_EPOCHS = 5
-LEARNING_RATE = 1e-4
-MODEL_SAVE_PATH = "models/deepfake_model.pth"
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# -----------------------------
-# DATA TRANSFORMS
-# -----------------------------
-train_transforms = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.RandomHorizontalFlip(),
-    transforms.ToTensor(),
-])
-
-val_transforms = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-])
-
-# -----------------------------
-# LOAD DATA
-# -----------------------------
-train_dataset = datasets.ImageFolder(os.path.join(DATA_DIR, "train"), transform=train_transforms)
-val_dataset = datasets.ImageFolder(os.path.join(DATA_DIR, "val"), transform=val_transforms)
-
-train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
-val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE)
-
-# -----------------------------
-# MODEL SETUP
-# -----------------------------
-# Use pretrained ResNet18
-model = models.resnet18(pretrained=True)
-
-# Replace the final layer with 2 classes (real/fake)
-num_features = model.fc.in_features
-model.fc = nn.Linear(num_features, 2)
-
-model = model.to(DEVICE)
-
-# Loss and optimizer
-criterion = nn.CrossEntropyLoss()
-optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
-
-# -----------------------------
-# TRAIN LOOP
-# -----------------------------
-for epoch in range(NUM_EPOCHS):
-    model.train()
-    running_loss = 0.0
-    correct = 0
-    total = 0
-    
-    for images, labels in train_loader:
-        images, labels = images.to(DEVICE), labels.to(DEVICE)
-        optimizer.zero_grad()
-        outputs = model(images)
-        loss = criterion(outputs, labels)
-        loss.backward()
-        optimizer.step()
-        
-        running_loss += loss.item()
-        _, predicted = torch.max(outputs, 1)
-        total += labels.size(0)
-        correct += (predicted == labels).sum().item()
-    
-    train_acc = correct / total
-    print(f"Epoch {epoch+1}/{NUM_EPOCHS} - Loss: {running_loss:.4f} - Train Accuracy: {train_acc:.4f}")
-    
-    # Validation
-    model.eval()
-    val_correct = 0
-    val_total = 0
-    with torch.no_grad():
-        for images, labels in val_loader:
-            images, labels = images.to(DEVICE), labels.to(DEVICE)
-            outputs = model(images)
-            _, predicted = torch.max(outputs, 1)
-            val_total += labels.size(0)
-            val_correct += (predicted == labels).sum().item()
-    val_acc = val_correct / val_total
-    print(f"Validation Accuracy: {val_acc:.4f}")
-
-# -----------------------------
-# SAVE MODEL
-# -----------------------------
-torch.save(model.state_dict(), MODEL_SAVE_PATH)
-print(f"Model saved to {MODEL_SAVE_PATH}")