""" Enhanced training script for Sequence VAE. Fixes: - Accepts both --latent_dim and --latent-dim arguments - KL annealing for stable training - Learning rate scheduling - Gradient clipping - Comprehensive metrics logging - Dual-path scoring calibration Usage: python ml/scripts/train_sequence_vae_enhanced.py \ --epochs 50 \ --beta 0.01 \ --latent-dim 16 \ --batch-size 64 \ --lr 0.001 \ --save-path artifacts/checkpoints/geometry/sequence_vae_enhanced.pt """ import argparse import json import sys import time from pathlib import Path import numpy as np import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, Dataset # Add repo root to path sys.path.append(str(Path(__file__).parent.parent.parent)) from ml.models.geometry.sequence_vae_enhanced import EnhancedSequenceVAE class GeometrySequenceDataset(Dataset): """Dataset for geometry sequences.""" def __init__(self, csv_path: str, seq_len: int = 80, n_features: int = 5): import pandas as pd self.seq_len = seq_len self.n_features = n_features # Load data df = pd.read_csv(csv_path) # Extract features feature_cols = ['twist_3m', 'versine_10m', 'versine_20m', 'unevenness_10m', 'cant'] alt_feature_cols = ['twist_3m_mm', 'versine_10m_mm', 'versine_20m_mm', 'unevenness_10m_mm', 'cant_mm'] if all(col in df.columns for col in feature_cols): self.data = df[feature_cols].values.astype(np.float32) elif all(col in df.columns for col in alt_feature_cols): self.data = df[alt_feature_cols].values.astype(np.float32) else: # Fallback: use first n_features numeric columns numeric_cols = df.select_dtypes(include=[np.number]).columns self.data = df[numeric_cols[:n_features]].values.astype(np.float32) # Normalize self.mean = self.data.mean(axis=0) self.std = self.data.std(axis=0) + 1e-8 self.data = (self.data - self.mean) / self.std # Create sequences self.sequences = [] stride = max(1, seq_len // 2) for i in range(0, len(self.data) - seq_len + 1, stride): # 50% overlap self.sequences.append(self.data[i:i + seq_len]) self.sequences = np.array(self.sequences) # Filter to normal only (if labels available) if 'label' in df.columns and len(self.sequences) > 0: # Keep only normal sequences (label == 0) normal_mask = df['label'].values[:len(self.sequences)] == 0 if normal_mask.sum() > 0: self.sequences = self.sequences[normal_mask] print(f"Loaded {len(self.sequences)} sequences from {csv_path}") def __len__(self): return len(self.sequences) def __getitem__(self, idx): return torch.tensor(self.sequences[idx], dtype=torch.float32) def train_enhanced_vae(args): """Train the enhanced Sequence VAE.""" print("=" * 70) print("TrackChain Enhanced Sequence VAE Training Pipeline") print("=" * 70) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # Print configuration print(f"\nConfiguration:") print(f" Epochs: {args.epochs}") print(f" Beta: {args.beta}") print(f" Latent dim: {args.latent_dim}") print(f" Batch size: {args.batch_size}") print(f" Learning rate: {args.lr}") print(f" KL annealing: {args.kl_annealing}") print(f" Annealing epochs: {args.annealing_epochs}") # Load data print(f"\n[1/6] Loading data from {args.data_path}...") # Find CSV file data_path = Path(args.data_path) csv_files = [] if data_path.exists(): if data_path.is_file() and data_path.suffix == '.csv': csv_files = [data_path] else: csv_files = list(data_path.glob('*.csv')) + list(data_path.glob('**/*.csv')) if not csv_files: # Generate synthetic data print(" No CSV found, generating synthetic data...") from ml.data.synthetic_geometry import SyntheticGeometryDataset dataset = SyntheticGeometryDataset(num_samples=3000, seq_len=args.seq_len) # Filter to normal only normal_indices = [i for i, (_, label) in enumerate(dataset) if label == 0] # Create sequences tensor sequences = [] for i in normal_indices: seq, _ = dataset[i] sequences.append(seq.numpy() if hasattr(seq, 'numpy') else np.asarray(seq)) sequences = np.array(sequences) else: csv_path = csv_files[0] print(f" Using: {csv_path}") dataset = GeometrySequenceDataset(str(csv_path), seq_len=args.seq_len, n_features=args.n_features) sequences = dataset.sequences if len(sequences) == 0: print(" CSV sequence count was 0, falling back to synthetic data...") from ml.data.synthetic_geometry import SyntheticGeometryDataset synth_dataset = SyntheticGeometryDataset(num_samples=3000, seq_len=args.seq_len) normal_indices = [i for i, (_, label) in enumerate(synth_dataset) if label == 0] sequences = np.array([synth_dataset[i][0].numpy() for i in normal_indices]) # Split into train and validation n_total = len(sequences) n_train = max(1, int(0.8 * n_total)) train_sequences = sequences[:n_train] val_sequences = sequences[n_train:] if n_train < n_total else sequences[:max(1, int(0.2 * n_total))] print(f" Train: {len(train_sequences)} sequences") print(f" Valid: {len(val_sequences)} sequences") # Create data loaders train_loader = DataLoader( torch.utils.data.TensorDataset(torch.tensor(train_sequences, dtype=torch.float32)), batch_size=min(args.batch_size, len(train_sequences)), shuffle=True, drop_last=(len(train_sequences) > args.batch_size), ) val_loader = DataLoader( torch.utils.data.TensorDataset(torch.tensor(val_sequences, dtype=torch.float32)), batch_size=min(args.batch_size, len(val_sequences)), shuffle=False, ) # Initialize model print(f"\n[2/6] Initializing Enhanced Sequence VAE...") model = EnhancedSequenceVAE( seq_len=args.seq_len, n_features=args.n_features, latent_dim=args.latent_dim, beta=args.beta, use_kl_annealing=args.kl_annealing, annealing_epochs=args.annealing_epochs, ).to(device) # Count parameters n_params = sum(p.numel() for p in model.parameters()) print(f" Model parameters: {n_params:,}") # Optimizer and scheduler optimizer = optim.AdamW( model.parameters(), lr=args.lr, weight_decay=args.weight_decay, ) scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=max(1, args.epochs), eta_min=args.lr * 0.01, ) # Training loop print(f"\n[3/6] Training for {args.epochs} epochs...") print("-" * 70) best_val_loss = float('inf') best_epoch = 0 training_history = { 'epochs': [], 'train_loss': [], 'train_recon': [], 'train_kl': [], 'val_loss': [], 'val_recon': [], 'val_kl': [], 'lr': [], 'beta': [], } start_time = time.time() for epoch in range(args.epochs): # Training model.train() train_total_loss = 0.0 train_recon_loss = 0.0 train_kl_loss = 0.0 current_beta = 0.0 for batch_idx, (batch_data,) in enumerate(train_loader): batch_data = batch_data.to(device) optimizer.zero_grad() # Forward pass recon_batch, mu, logvar = model(batch_data) # Compute loss with KL annealing losses = model.compute_loss(batch_data, recon_batch, mu, logvar, epoch) # Backward pass losses['total'].backward() # Gradient clipping torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=args.grad_clip) optimizer.step() train_total_loss += losses['total'].item() train_recon_loss += losses['recon'].item() train_kl_loss += losses['kl'].item() current_beta = float(losses['beta']) # Validation model.eval() val_total_loss = 0.0 val_recon_loss = 0.0 val_kl_loss = 0.0 with torch.no_grad(): for (batch_data,) in val_loader: batch_data = batch_data.to(device) recon_batch, mu, logvar = model(batch_data) losses = model.compute_loss(batch_data, recon_batch, mu, logvar, epoch) val_total_loss += losses['total'].item() val_recon_loss += losses['recon'].item() val_kl_loss += losses['kl'].item() # Compute averages n_train_batches = max(1, len(train_loader)) n_val_batches = max(1, len(val_loader)) avg_train_loss = train_total_loss / n_train_batches avg_train_recon = train_recon_loss / n_train_batches avg_train_kl = train_kl_loss / n_train_batches avg_val_loss = val_total_loss / n_val_batches avg_val_recon = val_recon_loss / n_val_batches avg_val_kl = val_kl_loss / n_val_batches # Log training_history['epochs'].append(epoch + 1) training_history['train_loss'].append(avg_train_loss) training_history['train_recon'].append(avg_train_recon) training_history['train_kl'].append(avg_train_kl) training_history['val_loss'].append(avg_val_loss) training_history['val_recon'].append(avg_val_recon) training_history['val_kl'].append(avg_val_kl) training_history['lr'].append(optimizer.param_groups[0]['lr']) training_history['beta'].append(current_beta) # Print progress if (epoch + 1) % 5 == 0 or epoch == 0 or epoch == args.epochs - 1: print(f"Epoch {epoch+1:02d}/{args.epochs} | " f"Train: {avg_train_loss:.4f} (R:{avg_train_recon:.4f} K:{avg_train_kl:.4f}) | " f"Val: {avg_val_loss:.4f} (R:{avg_val_recon:.4f} K:{avg_val_kl:.4f}) | " f"β:{current_beta:.4f} | " f"LR:{optimizer.param_groups[0]['lr']:.6f}") # Save best model and check early stopping if avg_val_loss < best_val_loss: best_val_loss = avg_val_loss best_epoch = epoch + 1 early_stopping_counter = 0 save_path = Path(args.save_path) save_path.parent.mkdir(parents=True, exist_ok=True) torch.save(model.state_dict(), save_path) else: early_stopping_counter += 1 if early_stopping_counter >= getattr(args, 'patience', 8): print(f"\n[INFO] Early stopping triggered at epoch {epoch+1} (patience={args.patience} reached, best epoch: {best_epoch})") break # Update learning rate scheduler.step() training_time = time.time() - start_time print(f"\n[4/6] Training completed in {training_time/60:.2f} minutes") print(f" Best epoch: {best_epoch}") print(f" Best val loss: {best_val_loss:.4f}") # Fit latent distribution for Mahalanobis scoring print(f"\n[5/6] Fitting latent distribution for dual-path scoring...") if Path(args.save_path).exists(): model.load_state_dict(torch.load(args.save_path, map_location=device)) model.eval() # Use validation sequences for fitting (model handles device mapping internally) val_tensor = torch.tensor(val_sequences, dtype=torch.float32) model.fit_latent_distribution(val_tensor) print(f" Latent mean shape: {model.latent_mean.shape}") print(f" Latent covariance inverse shape: {model.latent_cov_inv.shape}") # Calibrate thresholds with Extreme Value Theory (EVT) and P99 print(f"\n[6/6] Calibrating anomaly thresholds with Extreme Value Theory (EVT)...") # Compute reconstruction errors on validation set val_errors = [] val_mahalanobis = [] with torch.no_grad(): for i in range(0, len(val_sequences), args.batch_size): batch = torch.tensor(val_sequences[i:i + args.batch_size], dtype=torch.float32).to(device) for sample_idx in range(batch.size(0)): scores = model.compute_anomaly_score(batch[sample_idx]) val_errors.append(scores.get('recon_error', 0.0)) val_mahalanobis.append(scores.get('mahalanobis_dist', 0.0)) # Compute P99 thresholds p99_recon = float(np.percentile(val_errors, 99)) if val_errors else 1.0 p99_mahalanobis = float(np.percentile(val_mahalanobis, 99)) if val_mahalanobis else 1.0 ensemble_scores = [0.7 * e + 0.3 * m for e, m in zip(val_errors, val_mahalanobis)] if val_errors else [1.0] p99_ensemble = float(np.percentile(ensemble_scores, 99)) # SOTA: Compute Extreme Value Theory (EVT) Peaks-Over-Threshold evt_recon = model.fit_evt_threshold(val_errors, target_fpr=0.01) evt_ensemble = model.fit_evt_threshold(ensemble_scores, target_fpr=0.01) # Empirical False Positive Rate (FPR) Validation Check (Upgrade B) actual_fpr = float(np.mean([s > evt_ensemble['threshold'] for s in ensemble_scores])) print(f" P99 Ensemble Threshold: {p99_ensemble:.4f}") print(f" EVT Ensemble Threshold: {evt_ensemble['threshold']:.4f} (shape={evt_ensemble['shape']:.4f}, scale={evt_ensemble['scale']:.4f})") print(f" Empirical Validation FPR: {actual_fpr*100:.2f}% (Target: 1.00%, Tolerance <= 2.00%)") assert actual_fpr <= 0.02, f"FPR validation check failed: {actual_fpr:.4f} > 0.02" # Strict Normalization Guard Check (Upgrade C) model.threshold_evt = float(evt_ensemble['threshold']) test_prob_min = model.score_to_probability(0.0) test_prob_thresh = model.score_to_probability(evt_ensemble['threshold']) test_prob_extreme = model.score_to_probability(evt_ensemble['threshold'] * 5.0) assert 0.0 <= test_prob_min <= 0.20, f"Min score probability out of expected range: {test_prob_min}" assert abs(test_prob_thresh - 0.50) < 0.01, f"Threshold decision boundary not 0.50: {test_prob_thresh}" assert 0.90 <= test_prob_extreme <= 1.0, f"Extreme anomaly probability not near 1.0: {test_prob_extreme}" print(f" Probability Normalization: Verified (P(0)={test_prob_min:.3f}, P(Threshold)={test_prob_thresh:.3f}, P(Extreme)={test_prob_extreme:.3f})") # Save calibration manifest calibration = { 'threshold_evt': float(evt_ensemble['threshold']), 'threshold_p99': float(p99_ensemble), 'threshold_recon_evt': float(evt_recon['threshold']), 'threshold_recon_p99': float(p99_recon), 'threshold_mahalanobis_p99': float(p99_mahalanobis), 'evt_shape': float(evt_ensemble['shape']), 'evt_scale': float(evt_ensemble['scale']), 'evt_init_threshold': float(evt_ensemble['init_threshold']), 'steepness': 0.5, 'steepness_k': 2.0, 'target_fpr': 0.01, 'actual_fpr': actual_fpr, 'fpr_verified': True, 'normalization_guard_passed': True, 'model': 'sequence_vae_geometry_novel', 'val_samples': len(val_sequences), 'mean_recon_error': float(np.mean(val_errors)) if val_errors else 0.0, 'std_recon_error': float(np.std(val_errors)) if val_errors else 0.0, 'mean_mahalanobis': float(np.mean(val_mahalanobis)) if val_mahalanobis else 0.0, 'std_mahalanobis': float(np.std(val_mahalanobis)) if val_mahalanobis else 0.0, } calib_path = Path('artifacts/calibration/vae_calibration.json') calib_path.parent.mkdir(parents=True, exist_ok=True) with open(calib_path, 'w', encoding='utf-8') as f: json.dump(calibration, f, indent=2) # Mirror calibration to legacy path legacy_calib = Path('artifacts/calibration/sequence_vae_calibration.json') with open(legacy_calib, 'w', encoding='utf-8') as f: json.dump({ "method": "evt_peaks_over_threshold", "threshold_evt": float(evt_ensemble['threshold']), "threshold_p99": float(p99_ensemble), "steepness_k": 0.5, "percentile": 99.0, "is_fitted": True, "model": "sequence_vae_geometry_novel" }, f, indent=2) # Save training history history_path = Path(args.save_path).parent / 'training_history.json' history_path.parent.mkdir(parents=True, exist_ok=True) with open(history_path, 'w', encoding='utf-8') as f: json.dump(training_history, f, indent=2) # Save latent distribution latent_dist_path = Path(args.save_path).parent / 'latent_distribution.npz' np.savez( latent_dist_path, mean=model.latent_mean, cov_inv=model.latent_cov_inv, ) # Mirror checkpoints to canonical & legacy locations for seamless inference canonical_pt = Path('artifacts/checkpoints/geometry/sequence_vae.pt') canonical_pt.parent.mkdir(parents=True, exist_ok=True) torch.save(model.state_dict(), canonical_pt) weights_mirror = Path('ml/models/geometry/weights/sequence_vae.pt') weights_mirror.parent.mkdir(parents=True, exist_ok=True) torch.save(model.state_dict(), weights_mirror) print(f"\n" + "=" * 70) print("Enhanced Sequence VAE Training Complete!") print("=" * 70) print(f"Model saved to: {args.save_path}") print(f"Mirrored to: {canonical_pt}") print(f" {weights_mirror}") print(f"Calibration saved: {calib_path}") print(f"Training history: {history_path}") print(f"\nCalibration parameters:") print(f" P99 reconstruction error: {p99_recon:.4f}") print(f" P99 Mahalanobis distance: {p99_mahalanobis:.4f}") print(f" P99 ensemble score: {p99_ensemble:.4f}") print(f" Target FPR: 1%") return model, calibration if __name__ == "__main__": parser = argparse.ArgumentParser(description="Train Enhanced Sequence VAE") # Training parameters parser.add_argument('--epochs', type=int, default=50, help='Number of training epochs') parser.add_argument('--beta', type=float, default=0.01, help='Beta-VAE weight for KL divergence') parser.add_argument('--batch-size', '--batch_size', dest='batch_size', type=int, default=64, help='Batch size') parser.add_argument('--lr', type=float, default=0.001, help='Learning rate') parser.add_argument('--weight-decay', '--weight_decay', dest='weight_decay', type=float, default=1e-4, help='Weight decay') parser.add_argument('--grad-clip', '--grad_clip', dest='grad_clip', type=float, default=1.0, help='Gradient clipping norm') parser.add_argument('--patience', type=int, default=8, help='Early stopping patience') # Model parameters - ACCEPT BOTH HYPHEN AND UNDERSCORE FORMATS parser.add_argument('--latent-dim', '--latent_dim', dest='latent_dim', type=int, default=16, help='Latent space dimension') parser.add_argument('--seq-len', '--seq_len', dest='seq_len', type=int, default=80, help='Sequence length') parser.add_argument('--n-features', '--n_features', dest='n_features', type=int, default=5, help='Number of input features') # KL annealing parser.add_argument('--kl-annealing', '--kl_annealing', dest='kl_annealing', action='store_true', default=True, help='Use KL annealing') parser.add_argument('--no-kl-annealing', '--no_kl_annealing', dest='kl_annealing', action='store_false') parser.add_argument('--annealing-epochs', '--annealing_epochs', dest='annealing_epochs', type=int, default=10, help='Number of epochs for KL annealing') # Paths parser.add_argument('--save-path', '--save_path', dest='save_path', type=str, default='artifacts/checkpoints/geometry/sequence_vae_enhanced.pt') parser.add_argument('--data-path', '--data_path', dest='data_path', type=str, default='data/processed/normal_sequences/') args = parser.parse_args() train_enhanced_vae(args)