Add training script for all 3 models

experiments/train_baseline.py

@@ -0,0 +1,282 @@
+"""
+Training script for all three IDS models.
+Trains MLP, LSTM, and 1D-CNN on NSL-KDD with full evaluation.
+"""
+
+import os
+import sys
+import json
+import time
+import random
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import TensorDataset, DataLoader
+from sklearn.metrics import (classification_report, confusion_matrix,
+                             roc_auc_score, average_precision_score)
+
+# Add project root to path
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from models.mlp_baseline import MLP_IDS
+from models.lstm_model import LSTM_IDS
+from models.cnn1d_model import CNN1D_IDS
+from data.preprocess import load_nsl_kdd, preprocess, save_preprocessed, FEATURE_NAMES
+
+# ========================
+# Reproducibility
+# ========================
+SEED = 42
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+# ========================
+# Config
+# ========================
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+NUM_CLASSES = 2  # Binary classification
+EPOCHS = 50
+BATCH_SIZE = 256
+LR = 1e-3
+WEIGHT_DECAY = 1e-4
+RESULTS_DIR = 'results'
+MODELS_DIR = 'saved_models'
+
+
+def compute_class_weights(y_train):
+    """Compute inverse-frequency class weights."""
+    counts = np.bincount(y_train)
+    weights = 1.0 / counts.astype(np.float32)
+    weights = weights / weights.sum() * len(weights)  # Normalize
+    return torch.FloatTensor(weights).to(DEVICE)
+
+
+def train_one_epoch(model, loader, criterion, optimizer):
+    """Train for one epoch."""
+    model.train()
+    total_loss = 0
+    correct = 0
+    total = 0
+    
+    for X_batch, y_batch in loader:
+        X_batch, y_batch = X_batch.to(DEVICE), y_batch.to(DEVICE)
+        
+        optimizer.zero_grad()
+        outputs = model(X_batch)
+        loss = criterion(outputs, y_batch)
+        loss.backward()
+        optimizer.step()
+        
+        total_loss += loss.item() * len(y_batch)
+        preds = outputs.argmax(dim=1)
+        correct += (preds == y_batch).sum().item()
+        total += len(y_batch)
+    
+    return total_loss / total, correct / total
+
+
+@torch.no_grad()
+def evaluate(model, loader, criterion):
+    """Evaluate model on dataset."""
+    model.eval()
+    total_loss = 0
+    all_preds = []
+    all_probs = []
+    all_labels = []
+    
+    for X_batch, y_batch in loader:
+        X_batch, y_batch = X_batch.to(DEVICE), y_batch.to(DEVICE)
+        
+        outputs = model(X_batch)
+        loss = criterion(outputs, y_batch)
+        
+        total_loss += loss.item() * len(y_batch)
+        probs = torch.softmax(outputs, dim=1)
+        all_preds.append(outputs.argmax(dim=1).cpu().numpy())
+        all_probs.append(probs.cpu().numpy())
+        all_labels.append(y_batch.cpu().numpy())
+    
+    all_preds = np.concatenate(all_preds)
+    all_probs = np.concatenate(all_probs)
+    all_labels = np.concatenate(all_labels)
+    
+    avg_loss = total_loss / len(all_labels)
+    
+    return avg_loss, all_preds, all_probs, all_labels
+
+
+def full_evaluation(y_true, y_pred, y_probs, class_names):
+    """Compute all metrics."""
+    results = {}
+    
+    # Classification report
+    report = classification_report(y_true, y_pred, target_names=class_names, output_dict=True)
+    results['classification_report'] = report
+    
+    # ROC-AUC (binary)
+    if len(class_names) == 2:
+        results['roc_auc'] = roc_auc_score(y_true, y_probs[:, 1])
+        results['pr_auc'] = average_precision_score(y_true, y_probs[:, 1])
+    
+    # Confusion matrix
+    cm = confusion_matrix(y_true, y_pred)
+    results['confusion_matrix'] = cm.tolist()
+    
+    return results
+
+
+def train_model(model, model_name, X_train, y_train, X_test, y_test, class_names):
+    """Full training pipeline for one model."""
+    print(f"\n{'='*60}")
+    print(f"Training {model_name}")
+    print(f"{'='*60}")
+    print(f"Parameters: {model.count_parameters():,}")
+    print(f"Device: {DEVICE}")
+    
+    # Data loaders
+    train_ds = TensorDataset(torch.FloatTensor(X_train), torch.LongTensor(y_train))
+    test_ds = TensorDataset(torch.FloatTensor(X_test), torch.LongTensor(y_test))
+    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)
+    test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False)
+    
+    # Loss with class weights
+    class_weights = compute_class_weights(y_train)
+    criterion = nn.CrossEntropyLoss(weight=class_weights)
+    
+    # Optimizer
+    optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)
+    
+    # Training loop
+    model.to(DEVICE)
+    best_f1 = 0
+    history = {'train_loss': [], 'train_acc': [], 'test_loss': [], 'test_acc': []}
+    
+    start_time = time.time()
+    
+    for epoch in range(EPOCHS):
+        train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer)
+        test_loss, test_preds, test_probs, test_labels = evaluate(model, test_loader, criterion)
+        test_acc = (test_preds == test_labels).mean()
+        
+        scheduler.step(test_loss)
+        
+        history['train_loss'].append(train_loss)
+        history['train_acc'].append(train_acc)
+        history['test_loss'].append(test_loss)
+        history['test_acc'].append(test_acc)
+        
+        # Check for best model
+        report = classification_report(test_labels, test_preds, output_dict=True)
+        weighted_f1 = report['weighted avg']['f1-score']
+        
+        if weighted_f1 > best_f1:
+            best_f1 = weighted_f1
+            os.makedirs(MODELS_DIR, exist_ok=True)
+            torch.save(model.state_dict(), os.path.join(MODELS_DIR, f'{model_name}_best.pt'))
+        
+        if (epoch + 1) % 10 == 0 or epoch == 0:
+            print(f"  Epoch {epoch+1:3d}/{EPOCHS} | "
+                  f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} | "
+                  f"Test Loss: {test_loss:.4f} Acc: {test_acc:.4f} F1: {weighted_f1:.4f}")
+    
+    train_time = time.time() - start_time
+    print(f"\n  Training time: {train_time:.1f}s")
+    
+    # Load best model and final evaluation
+    model.load_state_dict(torch.load(os.path.join(MODELS_DIR, f'{model_name}_best.pt'),
+                                      weights_only=True))
+    _, final_preds, final_probs, final_labels = evaluate(model, test_loader, criterion)
+    
+    results = full_evaluation(final_labels, final_preds, final_probs, class_names)
+    results['training_time'] = train_time
+    results['best_weighted_f1'] = best_f1
+    results['history'] = history
+    results['parameters'] = model.count_parameters()
+    
+    # Print final results
+    print(f"\n  Final Results ({model_name}):")
+    print(f"  {'='*50}")
+    print(classification_report(final_labels, final_preds, target_names=class_names))
+    
+    if 'roc_auc' in results:
+        print(f"  ROC-AUC: {results['roc_auc']:.4f}")
+        print(f"  PR-AUC:  {results['pr_auc']:.4f}")
+    
+    print(f"  Confusion Matrix:\n{confusion_matrix(final_labels, final_preds)}")
+    
+    return model, results
+
+
+def main():
+    # ========================
+    # Data
+    # ========================
+    df_train, df_test = load_nsl_kdd()
+    X_train, X_test, y_train, y_test, le, scaler, class_names = preprocess(
+        df_train, df_test, binary=True
+    )
+    save_preprocessed(X_train, X_test, y_train, y_test, le, scaler, class_names)
+    
+    # ========================
+    # Train all models
+    # ========================
+    all_results = {}
+    
+    # 1. MLP Baseline
+    mlp = MLP_IDS(in_dim=41, num_classes=NUM_CLASSES)
+    mlp, mlp_results = train_model(mlp, 'mlp', X_train, y_train, X_test, y_test, class_names)
+    all_results['mlp'] = mlp_results
+    
+    # 2. LSTM
+    lstm = LSTM_IDS(in_dim=41, num_classes=NUM_CLASSES)
+    lstm, lstm_results = train_model(lstm, 'lstm', X_train, y_train, X_test, y_test, class_names)
+    all_results['lstm'] = lstm_results
+    
+    # 3. 1D-CNN
+    cnn = CNN1D_IDS(in_dim=41, num_classes=NUM_CLASSES)
+    cnn, cnn_results = train_model(cnn, 'cnn1d', X_train, y_train, X_test, y_test, class_names)
+    all_results['cnn1d'] = cnn_results
+    
+    # ========================
+    # Save results
+    # ========================
+    os.makedirs(RESULTS_DIR, exist_ok=True)
+    
+    def convert(o):
+        if isinstance(o, np.floating): return float(o)
+        if isinstance(o, np.integer): return int(o)
+        if isinstance(o, np.ndarray): return o.tolist()
+        return o
+    
+    with open(os.path.join(RESULTS_DIR, 'training_results.json'), 'w') as f:
+        json.dump(all_results, f, indent=2, default=convert)
+    
+    # ========================
+    # Summary comparison
+    # ========================
+    print("\n" + "="*60)
+    print("MODEL COMPARISON SUMMARY")
+    print("="*60)
+    print(f"{'Model':<10} {'Params':>8} {'Accuracy':>10} {'W-F1':>8} {'ROC-AUC':>9} {'PR-AUC':>8} {'Time':>8}")
+    print("-"*60)
+    
+    for name, res in all_results.items():
+        acc = res['classification_report']['accuracy']
+        wf1 = res['best_weighted_f1']
+        roc = res.get('roc_auc', 0)
+        pr = res.get('pr_auc', 0)
+        t = res['training_time']
+        p = res['parameters']
+        print(f"{name:<10} {p:>8,} {acc:>10.4f} {wf1:>8.4f} {roc:>9.4f} {pr:>8.4f} {t:>7.1f}s")
+    
+    print("\nAll models trained successfully!")
+    print(f"Results saved to {RESULTS_DIR}/training_results.json")
+    print(f"Models saved to {MODELS_DIR}/")
+
+
+if __name__ == '__main__':
+    main()