longtrain.py

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader, WeightedRandomSampler, Dataset
import os
from pathlib import Path
from tqdm import tqdm
from sklearn.metrics import roc_auc_score, confusion_matrix
import numpy as np
from PIL import Image

# --- Configuration ---
TARGET_ARCH = "32bit" 
DATA_DIR = Path(f"data_processed_filtered/{TARGET_ARCH}") 
CHALLENGE_DIR = Path("data_separated/32bit/challenge") 

# OVERNIGHT SETTINGS
BATCH_SIZE = 8        # Reduced to fit B4 + 512px images in VRAM (Float32)
EPOCHS = 50            # Long run to allow deep convergence
LEARNING_RATE = 0.0002 # Lower LR for fine-tuning a large model
IMG_SIZE = 512         # High Resolution

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MEMMAP_THRESHOLD = 100 * 1024 * 1024 
Image.MAX_IMAGE_PIXELS = None 

# --- 1. AUGMENTATION: NOISE INJECTION ---
class AddGaussianNoise(object):
    def __init__(self, mean=0., std=0.1):
        self.std = std
        self.mean = mean
        
    def __call__(self, tensor):
        return tensor + torch.randn(tensor.size()) * self.std + self.mean

# --- 2. LOSS: FOCAL LOSS ---
class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.ce = nn.CrossEntropyLoss(reduction='none')

    def forward(self, inputs, targets):
        ce_loss = self.ce(inputs, targets)
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
        return focal_loss.mean()

# --- 3. DATA UTILS ---
def get_width(file_size_kb):
    if file_size_kb < 10: return 32
    if file_size_kb < 30: return 64
    if file_size_kb < 60: return 128
    if file_size_kb < 100: return 256
    if file_size_kb < 200: return 384
    if file_size_kb < 1000: return 512
    if file_size_kb < 1500: return 1024
    return 2048

class ChallengeDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = Path(root_dir)
        self.transform = transform
        self.file_paths = [p for p in self.root_dir.rglob("*") if p.is_file()]
        self.class_map = {'goodware': 0, 'malware': 1, 'benign': 0, 'malicious': 1}

    def __len__(self):
        return len(self.file_paths)
    
    def get_label(self, path):
        for part in path.parts:
            if part.lower() in self.class_map:
                return self.class_map[part.lower()]
        return -1

    def __getitem__(self, idx):
        file_path = self.file_paths[idx]
        label = self.get_label(file_path)
        try:
            file_stat = os.stat(file_path)
            file_size = file_stat.st_size
            if file_size == 0:
                img = Image.new('L', (IMG_SIZE, IMG_SIZE), 0)
            else:
                width = get_width(file_size / 1024)
                height = file_size // width
                if height == 0: height = 1
                
                if file_size > MEMMAP_THRESHOLD:
                    img_array = np.memmap(file_path, dtype=np.uint8, mode='r', shape=(height, width))
                    img = Image.fromarray(img_array, 'L')
                    del img_array
                else:
                    with open(file_path, 'rb') as f:
                        data = np.frombuffer(f.read(height * width), dtype=np.uint8)
                    if data.size != height * width:
                        img = Image.new('L', (IMG_SIZE, IMG_SIZE), 0)
                    else:
                        img = Image.fromarray(data.reshape((height, width)), 'L')

            if self.transform: img = self.transform(img)
            return img, label
        except:
            return torch.zeros((3, IMG_SIZE, IMG_SIZE)), -1

def get_loaders():
    ts_train = transforms.Compose([
        transforms.Grayscale(3), 
        transforms.Resize((IMG_SIZE, IMG_SIZE)),
        transforms.RandomHorizontalFlip(), 
        transforms.RandomVerticalFlip(),
        # Add rotation to force model to learn local features (headers) regardless of orientation
        transforms.RandomRotation(15), 
        transforms.ToTensor(), 
        AddGaussianNoise(0., 0.05), 
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # ImageNet Stats for B4
    ])
    ts_val = transforms.Compose([
        transforms.Grayscale(3), 
        transforms.Resize((IMG_SIZE, IMG_SIZE)),
        transforms.ToTensor(), 
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    image_datasets = {x: datasets.ImageFolder(os.path.join(DATA_DIR, x), ts_train if x=='train' else ts_val)
                      for x in ['train', 'val']}
    
    targets = image_datasets['train'].targets
    class_counts = np.bincount(targets)
    weights = 1. / class_counts
    samples_weights = [weights[t] for t in targets]
    sampler = WeightedRandomSampler(samples_weights, len(samples_weights))

    train_loader = DataLoader(image_datasets['train'], batch_size=BATCH_SIZE, sampler=sampler, num_workers=4, pin_memory=True)
    val_loader = DataLoader(image_datasets['val'], batch_size=BATCH_SIZE, shuffle=False, num_workers=4, pin_memory=True)
    
    challenge_ds = ChallengeDataset(CHALLENGE_DIR, transform=ts_val)
    challenge_loader = DataLoader(challenge_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=4, pin_memory=True)

    return train_loader, val_loader, challenge_loader, len(image_datasets['train'])

def calculate_metrics(y_true, y_probs, y_preds):
    valid = [i for i, x in enumerate(y_true) if x != -1]
    if not valid: return 0, 0, 0, 0
    y_true = [y_true[i] for i in valid]
    y_probs = [y_probs[i] for i in valid]
    y_preds = [y_preds[i] for i in valid]

    try: auc = roc_auc_score(y_true, y_probs)
    except: auc = 0.5
    tn, fp, fn, tp = confusion_matrix(y_true, y_preds, labels=[0, 1]).ravel()
    fpr = fp / (fp + tn) if (fp + tn) > 0 else 0.0
    fnr = fn / (fn + tp) if (fn + tp) > 0 else 0.0
    acc = (tp + tn) / (tp + tn + fp + fn)
    return acc, auc, fpr, fnr

def main():
    train_loader, val_loader, challenge_loader, train_size = get_loaders()
    print(f"Training EfficientNet-B4 (Overnight) @ {IMG_SIZE}x{IMG_SIZE}")
    print(f"Batches: {len(train_loader)} per epoch")
    
    # Load Heavy Model
    model = models.efficientnet_b4(weights='IMAGENET1K_V1')
    num_ftrs = model.classifier[1].in_features
    model.classifier[1] = nn.Linear(num_ftrs, 2)
    model = model.to(DEVICE)
    
    criterion = FocalLoss(alpha=0.25, gamma=2.0)
    # AdamW handles weight decay better for long runs
    optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=0.01)
    # Cosine Annealing to 0 ensures we settle into the best local minima
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)

    best_val_acc = 0.0
    best_chal_auc = 0.0

    for epoch in range(EPOCHS):
        print(f'\nEpoch {epoch+1}/{EPOCHS}')
        print('-' * 40)

        model.train()
        run_loss, run_corr = 0.0, 0
        loop = tqdm(train_loader, desc="Train", leave=False)
        
        for inputs, labels in loop:
            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            run_loss += loss.item() * inputs.size(0)
            _, preds = torch.max(outputs, 1)
            run_corr += torch.sum(preds == labels.data)
            loop.set_postfix(loss=loss.item())
        
        # Step scheduler after epoch
        scheduler.step()
        epoch_loss = run_loss / train_size
        epoch_acc = run_corr.double() / train_size
        print(f"| Train      | Loss: {epoch_loss:.4f} | Acc: {epoch_acc:.4f}")

        model.eval()
        # Validate
        for phase, loader in [('Validation', val_loader), ('Challenge', challenge_loader)]:
            if len(loader.dataset) == 0: continue
            all_labels, all_probs, all_preds = [], [], []
            
            with torch.no_grad():
                for inputs, labels in tqdm(loader, desc=phase, leave=False):
                    inputs = inputs.to(DEVICE)
                    outputs = model(inputs)
                    probs = torch.softmax(outputs, dim=1)
                    _, preds = torch.max(outputs, 1)
                    
                    all_labels.extend(labels.cpu().numpy())
                    all_probs.extend(probs[:, 1].cpu().numpy())
                    all_preds.extend(preds.cpu().numpy())

            acc, auc, fpr, fnr = calculate_metrics(all_labels, all_probs, all_preds)
            print(f"| {phase:<10} | AUC: {auc:.4f} | Acc: {acc:.4f} | FPR: {fpr:.4f} | FNR: {fnr:.4f}")

            # Save checkpoints
            if phase == 'Validation' and acc > best_val_acc:
                best_val_acc = acc
                torch.save(model.state_dict(), f"{TARGET_ARCH}_overnight_best_val.pth")
            
            # Also save best Challenge AUC so you don't miss the peak if it overfits later
            if phase == 'Challenge' and auc > best_chal_auc:
                best_chal_auc = auc
                torch.save(model.state_dict(), f"{TARGET_ARCH}_overnight_best_challenge.pth")

    torch.save(model.state_dict(), f"{TARGET_ARCH}_overnight_final.pth")
    print("\nTraining Complete.")

if __name__ == '__main__':
    import multiprocessing
    multiprocessing.freeze_support()
    main()