""" train_balanced.py ================== Trains the two balanced-training ablation variants on top of the splits produced by balance_data.py: --mode balanced_recon Same unsupervised reconstruction-error training as train.py's normal-only model -- MSE loss, no labels used -- but fed the balanced normal+attack training set instead of normal-only. Tests whether reconstruction error still separates the classes when the model is trained to reconstruct BOTH well. --mode balanced_supervised Labels are used directly. The encoder's bottleneck vector feeds a small classifier head trained with binary cross-entropy, combined with the reconstruction loss (weighted). This is structurally closest to the Elsayed et al. (LSTM-AE + OC-SVM on the latent representation) precedent already cited in Literature.tex, except the classifier head is trained jointly rather than as a separate downstream SVM step. Both variants share the same encoder/decoder architecture, hidden size, window size, and seed as the canonical normal-only model in train.py -- only the training data and (for balanced_supervised) the loss function differ. Run balance_data.py first to generate the required *_balanced_{run_id}.npy arrays. Usage ----- python balance_data.py --dataset csic2010 --window 5 python train_balanced.py --dataset csic2010 --mode balanced_recon python train_balanced.py --dataset csic2010 --mode balanced_supervised Author : K.A.D.S.D. Kandanaarachchi (2020/ICT/19) Project: Detecting Anomalous REST API Traffic -- IT4216 (balanced- training ablation) """ import argparse import json import logging import random from pathlib import Path import numpy as np import torch import torch.nn as nn from torch.utils.data import DataLoader, TensorDataset from model import ( LSTMAutoencoder, build_model_cicids2018, build_model_csic2010, build_model_unsw, ) logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s", datefmt="%H:%M:%S", ) log = logging.getLogger(__name__) def set_seed(seed=42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.backends.cudnn.deterministic = True class LSTMAutoencoderWithHead(nn.Module): """ Wraps the existing LSTMAutoencoder (encoder+decoder unchanged) and adds a small linear classifier head on top of the encoder's bottleneck vector, for the balanced_supervised variant. Kept as a separate wrapper class -- rather than modifying model.py's LSTMAutoencoder directly -- so the canonical normal-only model and all existing scripts (evaluate.py, active_learning.py, tui.py) are completely unaffected by this ablation study. """ def __init__(self, base_model: LSTMAutoencoder, hidden_size: int = 64): super().__init__() self.base = base_model self.classifier = nn.Linear(hidden_size, 1) def forward(self, x): if self.base.input_mode == "embedding": x_emb = self.base.embedding(x) else: x_emb = x context = self.base.encoder(x_emb) # bottleneck vector reconstruction = self.base.decoder(context) logit = self.classifier(context).squeeze(-1) return reconstruction, logit, x_emb def reconstruction_error(self, x): with torch.no_grad(): recon, _, x_emb = self.forward(x) error = ((recon - x_emb) ** 2).mean(dim=(1, 2)) return error def predict_proba(self, x): with torch.no_grad(): _, logit, _ = self.forward(x) return torch.sigmoid(logit) def _to_tensor(X: np.ndarray, dataset: str) -> torch.Tensor: if dataset == "csic2010": return torch.tensor(X, dtype=torch.long) X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0) return torch.tensor(X, dtype=torch.float32) def train_balanced_recon( model, X_train, dataset, run_id, epochs, batch_size, lr, device ): """ Reconstruction-only training on the balanced set -- identical training loop to train.py's train(), just fed balanced data instead of normal-only data. No labels used. """ model = model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=lr) criterion = nn.MSELoss() split = int(len(X_train) * 0.9) X_tr, X_val = X_train[:split], X_train[split:] tr_tensor, val_tensor = _to_tensor(X_tr, dataset), _to_tensor(X_val, dataset) tr_loader = DataLoader(TensorDataset(tr_tensor), batch_size=batch_size, shuffle=True) val_loader = DataLoader(TensorDataset(val_tensor), batch_size=batch_size) train_losses, val_losses = [], [] best_val = float("inf") for epoch in range(1, epochs + 1): model.train() epoch_loss = 0.0 for (batch,) in tr_loader: batch = batch.to(device) optimizer.zero_grad() recon = model(batch) target = model.embedding(batch).detach() if dataset == "csic2010" else batch loss = criterion(recon, target) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() epoch_loss += loss.item() * len(batch) epoch_loss /= len(X_tr) model.eval() val_loss = 0.0 with torch.no_grad(): for (batch,) in val_loader: batch = batch.to(device) recon = model(batch) target = model.embedding(batch).detach() if dataset == "csic2010" else batch val_loss += criterion(recon, target).item() * len(batch) val_loss /= len(X_val) train_losses.append(epoch_loss) val_losses.append(val_loss) log.info("Epoch %02d/%02d train=%.6f val=%.6f", epoch, epochs, epoch_loss, val_loss) if val_loss < best_val - 1e-6: best_val = val_loss torch.save(model.state_dict(), f"models/best_{run_id}_balanced_recon.pt") return train_losses, val_losses def train_balanced_supervised( model, X_train, y_train, dataset, run_id, epochs, batch_size, lr, device, recon_weight: float = 0.5, ): """ Joint reconstruction + classification training. Loss = recon_weight * MSE(reconstruction, input) + (1 - recon_weight) * BCE(classifier_logit, label) recon_weight=0.5 is a starting point, not a tuned value -- treat it as a hyperparameter to sweep if time allows. At recon_weight=0 this degenerates to a pure sequence classifier (no autoencoding objective at all); at 1.0 it's identical to balanced_recon and the classifier head is trained but never influences the shared encoder weights via backprop on the classification loss. """ model = model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=lr) mse = nn.MSELoss() bce = nn.BCEWithLogitsLoss() split = int(len(X_train) * 0.9) X_tr, X_val = X_train[:split], X_train[split:] y_tr, y_val = y_train[:split], y_train[split:] tr_tensor, val_tensor = _to_tensor(X_tr, dataset), _to_tensor(X_val, dataset) y_tr_t = torch.tensor(y_tr, dtype=torch.float32) y_val_t = torch.tensor(y_val, dtype=torch.float32) tr_loader = DataLoader(TensorDataset(tr_tensor, y_tr_t), batch_size=batch_size, shuffle=True) val_loader = DataLoader(TensorDataset(val_tensor, y_val_t), batch_size=batch_size) train_losses, val_losses = [], [] best_val = float("inf") for epoch in range(1, epochs + 1): model.train() epoch_loss = 0.0 for batch, labels in tr_loader: batch, labels = batch.to(device), labels.to(device) optimizer.zero_grad() recon, logit, x_emb = model(batch) loss = recon_weight * mse(recon, x_emb) + (1 - recon_weight) * bce(logit, labels) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() epoch_loss += loss.item() * len(batch) epoch_loss /= len(X_tr) model.eval() val_loss = 0.0 correct = 0 with torch.no_grad(): for batch, labels in val_loader: batch, labels = batch.to(device), labels.to(device) recon, logit, x_emb = model(batch) loss = recon_weight * mse(recon, x_emb) + (1 - recon_weight) * bce(logit, labels) val_loss += loss.item() * len(batch) pred = (torch.sigmoid(logit) > 0.5).float() correct += (pred == labels).sum().item() val_loss /= len(X_val) val_acc = correct / len(X_val) train_losses.append(epoch_loss) val_losses.append(val_loss) log.info( "Epoch %02d/%02d train=%.6f val=%.6f val_acc=%.4f", epoch, epochs, epoch_loss, val_loss, val_acc, ) if val_loss < best_val - 1e-6: best_val = val_loss torch.save(model.state_dict(), f"models/best_{run_id}_balanced_supervised.pt") return train_losses, val_losses def main(): set_seed(42) parser = argparse.ArgumentParser() parser.add_argument("--dataset", required=True, choices=["csic2010", "cicids2018", "unsw"]) parser.add_argument("--mode", required=True, choices=["balanced_recon", "balanced_supervised"]) parser.add_argument("--window", type=int, default=5) parser.add_argument("--epochs", type=int, default=30) parser.add_argument("--batch_size", type=int, default=256) parser.add_argument("--lr", type=float, default=1e-3) parser.add_argument("--recon_weight", type=float, default=0.5) args = parser.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" data_dir = Path("data/processed") model_dir = Path("models") model_dir.mkdir(exist_ok=True) run_id = f"{args.dataset}_w{args.window}" X_train = np.load(data_dir / f"X_train_balanced_{run_id}.npy") y_train = np.load(data_dir / f"y_train_balanced_{run_id}.npy") log.info( "Loaded balanced training set: %s (normal=%d attack=%d)", X_train.shape, int((y_train == 0).sum()), int((y_train == 1).sum()), ) if args.dataset == "csic2010": vocab_data = json.load(open(data_dir / f"vocab_{run_id}.json")) vocab = vocab_data.get("vocab", vocab_data) base_model = build_model_csic2010(vocab_size=len(vocab), seq_len=args.window) elif args.dataset == "cicids2018": base_model = build_model_cicids2018(n_features=X_train.shape[2], seq_len=args.window) else: base_model = build_model_unsw(n_features=X_train.shape[2], seq_len=args.window) if args.mode == "balanced_recon": train_losses, val_losses = train_balanced_recon( base_model, X_train, args.dataset, run_id, args.epochs, args.batch_size, args.lr, device, ) suffix = "balanced_recon" else: model = LSTMAutoencoderWithHead(base_model, hidden_size=64) train_losses, val_losses = train_balanced_supervised( model, X_train, y_train, args.dataset, run_id, args.epochs, args.batch_size, args.lr, device, recon_weight=args.recon_weight, ) suffix = "balanced_supervised" history = { "dataset": args.dataset, "window": args.window, "mode": args.mode, "epochs": list(range(1, len(train_losses) + 1)), "train_loss": train_losses, "val_loss": val_losses, } history_path = model_dir / f"history_{run_id}_{suffix}.json" with open(history_path, "w") as f: json.dump(history, f, indent=2) log.info("History saved -> %s", history_path) log.info("Best model saved -> models/best_%s_%s.pt", run_id, suffix) log.info( "NEXT: evaluate against data/processed/X_test_balanced_%s.npy " "(NOT the canonical X_test_%s.npy -- different, smaller test " "population; see balance_data.py's printed summary).", run_id, run_id, ) if __name__ == "__main__": main()