Buckets:
| """ | |
| train.py | |
| Autonomous training script for baseline CNN-LSTM EEG Seizure Detection pipeline. | |
| Supports checkpointing, early stopping, resumption, mixed-precision, and class-weighted cross-entropy. | |
| """ | |
| import os | |
| import sys | |
| import json | |
| import time | |
| import argparse | |
| import random | |
| import logging | |
| from datetime import datetime | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| import torch.nn as nn | |
| from torch.amp import autocast, GradScaler | |
| from sklearn.metrics import ( | |
| precision_score, recall_score, f1_score, | |
| precision_recall_curve, auc, roc_auc_score, | |
| confusion_matrix | |
| ) | |
| from model import build_model_from_config | |
| from dataset import get_data_loaders | |
| def setup_logger(log_dir): | |
| os.makedirs(log_dir, exist_ok=True) | |
| log_file = os.path.join(log_dir, "training.log") | |
| logger = logging.getLogger("EEG_Trainer") | |
| logger.setLevel(logging.INFO) | |
| logger.handlers.clear() | |
| fh = logging.FileHandler(log_file, mode="a", encoding="utf-8") | |
| fh.setLevel(logging.INFO) | |
| ch = logging.StreamHandler(sys.stdout) | |
| ch.setLevel(logging.INFO) | |
| fmt = logging.Formatter("[%(asctime)s][%(levelname)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S") | |
| fh.setFormatter(fmt) | |
| ch.setFormatter(fmt) | |
| logger.addHandler(fh) | |
| logger.addHandler(ch) | |
| return logger | |
| def set_seed(seed=42): | |
| random.seed(seed) | |
| np.random.seed(seed) | |
| torch.manual_seed(seed) | |
| if torch.cuda.is_available(): | |
| torch.cuda.manual_seed_all(seed) | |
| torch.backends.cudnn.deterministic = True | |
| torch.backends.cudnn.benchmark = False | |
| def compute_metrics(y_true, y_probs, y_preds): | |
| """ | |
| Compute binary classification metrics: | |
| Precision, Recall/Sensitivity, Specificity, F1, PR-AUC, ROC-AUC, CM. | |
| """ | |
| cm = confusion_matrix(y_true, y_preds, labels=[0, 1]) | |
| tn, fp, fn, tp = cm.ravel() if cm.shape == (2, 2) else (0, 0, 0, 0) | |
| precision = float(precision_score(y_true, y_preds, pos_label=1, zero_division=0)) | |
| sensitivity = float(recall_score(y_true, y_preds, pos_label=1, zero_division=0)) | |
| specificity = float(tn / (tn + fp)) if (tn + fp) > 0 else 0.0 | |
| f1 = float(f1_score(y_true, y_preds, pos_label=1, zero_division=0)) | |
| # PR-AUC | |
| try: | |
| p_curve, r_curve, _ = precision_recall_curve(y_true, y_probs) | |
| pr_auc = float(auc(r_curve, p_curve)) | |
| except Exception: | |
| pr_auc = 0.0 | |
| # ROC-AUC | |
| try: | |
| roc_auc = float(roc_auc_score(y_true, y_probs)) if len(np.unique(y_true)) > 1 else 0.5 | |
| except Exception: | |
| roc_auc = 0.5 | |
| return { | |
| "precision": precision, | |
| "sensitivity": sensitivity, | |
| "specificity": specificity, | |
| "f1": f1, | |
| "pr_auc": pr_auc, | |
| "roc_auc": roc_auc, | |
| "tn": int(tn), | |
| "fp": int(fp), | |
| "fn": int(fn), | |
| "tp": int(tp) | |
| } | |
| def evaluate_split(model, data_loader, criterion, device, use_amp=False, threshold=0.5, **kwargs): | |
| """ | |
| Evaluates model on a dataset split strictly in Float32 precision (no AMP), | |
| ensuring numerical parity and exact threshold consistency with evaluate.py. | |
| Args: | |
| model: PyTorch model. | |
| data_loader: DataLoader for evaluation. | |
| criterion: Loss function. | |
| device: torch.device (cuda or cpu). | |
| use_amp: Legacy boolean argument preserved for backwards compatibility. | |
| Validation inference is strictly executed in Float32 regardless | |
| of this parameter. | |
| threshold: Classification decision threshold for binary predictions (default: 0.5). | |
| **kwargs: Absorbs legacy arguments for backwards compatibility. | |
| """ | |
| model.eval() | |
| total_loss = 0.0 | |
| all_y_true = [] | |
| all_y_probs = [] | |
| all_y_preds = [] | |
| with torch.no_grad(): | |
| for x, y, _ in data_loader: | |
| x = x.to(device, non_blocking=True) | |
| y = y.to(device, non_blocking=True) | |
| # Explicitly execute evaluation forward pass in pure Float32 (no autocast) | |
| logits = model(x.float()) | |
| loss = criterion(logits, y) | |
| total_loss += loss.item() * len(y) | |
| probs = torch.softmax(logits.float(), dim=1)[:, 1].detach().cpu().numpy().astype(np.float32) | |
| preds = (probs >= threshold).astype(np.int64) | |
| all_y_true.extend(y.cpu().numpy().astype(np.int64)) | |
| all_y_probs.extend(probs) | |
| all_y_preds.extend(preds) | |
| avg_loss = total_loss / len(data_loader.dataset) | |
| all_y_true = np.array(all_y_true, dtype=np.int64) | |
| all_y_probs = np.array(all_y_probs, dtype=np.float32) | |
| all_y_preds = np.array(all_y_preds, dtype=np.int64) | |
| metrics = compute_metrics(all_y_true, all_y_probs, all_y_preds) | |
| metrics["loss"] = float(avg_loss) | |
| metrics["threshold"] = float(threshold) | |
| return metrics | |
| def train_pipeline(config_path="configs/baseline_config.json", resume=False): | |
| # Load configuration | |
| with open(config_path, "r") as f: | |
| config = json.load(f) | |
| train_cfg = config["training"] | |
| hw_cfg = config.get("hardware", {}) | |
| log_dir = train_cfg.get("log_dir", "logs") | |
| ckpt_dir = train_cfg.get("checkpoint_dir", "checkpoints") | |
| os.makedirs(ckpt_dir, exist_ok=True) | |
| os.makedirs(log_dir, exist_ok=True) | |
| logger = setup_logger(log_dir) | |
| logger.info("=" * 70) | |
| logger.info(f"STARTING CNN-LSTM TRAINING PIPELINE") | |
| logger.info(f"Experiment: {config.get('experiment_name')}") | |
| logger.info(f"Config path: {config_path}") | |
| # Set seed | |
| seed = config.get("random_seed", 42) | |
| set_seed(seed) | |
| logger.info(f"Random seed initialized to {seed}") | |
| # Hardware & Device | |
| use_cuda = hw_cfg.get("device", "cuda") == "cuda" and torch.cuda.is_available() | |
| device = torch.device("cuda" if use_cuda else "cpu") | |
| use_amp = hw_cfg.get("mixed_precision", True) and use_cuda | |
| logger.info(f"Compute Device: {device} (CUDA Available: {torch.cuda.is_available()})") | |
| if use_cuda: | |
| logger.info(f"GPU: {torch.cuda.get_device_name(0)} | VRAM: {torch.cuda.get_device_properties(0).total_memory / (1024**3):.2f} GB") | |
| logger.info(f"Mixed Precision (AMP): {use_amp}") | |
| # Prepare DataLoaders | |
| logger.info("Initializing DataLoaders with memory mapping...") | |
| data_dict = get_data_loaders(config) | |
| train_loader = data_dict["train_loader"] | |
| val_loader = data_dict["val_loader"] | |
| class_weights_np = data_dict["class_weights"] | |
| # Class-weighted loss with optional label smoothing | |
| label_smoothing = train_cfg.get("label_smoothing", 0.0) | |
| if train_cfg.get("use_class_weights", True): | |
| class_weights_t = torch.tensor(class_weights_np, dtype=torch.float32).to(device) | |
| criterion = nn.CrossEntropyLoss(weight=class_weights_t, label_smoothing=label_smoothing) | |
| logger.info(f"Using Class-Weighted CrossEntropyLoss: {class_weights_np.tolist()} (label_smoothing={label_smoothing})") | |
| else: | |
| criterion = nn.CrossEntropyLoss(label_smoothing=label_smoothing) | |
| logger.info(f"Using standard unweighted CrossEntropyLoss (label_smoothing={label_smoothing})") | |
| # Build Model | |
| model = build_model_from_config(config).to(device) | |
| total_params = sum(p.numel() for p in model.parameters() if p.requires_grad) | |
| logger.info(f"Model instantiated. Trainable parameters: {total_params:,}") | |
| logger.info( | |
| f"Regularization: LSTM dropout={getattr(model, 'lstm_dropout', model.lstm.dropout)}, " | |
| f"Classifier dropout={getattr(model, 'classifier_dropout', 'N/A')}, " | |
| f"Feat dropout={getattr(model, 'feat_dropout_rate', 0.0)}" | |
| ) | |
| # Smoke test forward pass | |
| logger.info("Executing smoke test forward pass...") | |
| test_x = torch.randn(2, 18, 2560).to(device) | |
| with autocast(device_type=device.type, enabled=use_amp): | |
| test_out = model(test_x) | |
| assert test_out.shape == (2, 2), f"Smoke test failed! Output shape: {test_out.shape}" | |
| logger.info(f"Smoke test passed! Logits shape: {test_out.shape}") | |
| # Optimizer & Scheduler | |
| lr = train_cfg.get("learning_rate", 5e-4) | |
| weight_decay = train_cfg.get("weight_decay", 1e-4) | |
| optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay) | |
| logger.info(f"Optimizer: AdamW (lr={lr}, weight_decay={weight_decay})") | |
| num_epochs = train_cfg.get("num_epochs", 15) | |
| scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=1e-6) | |
| scaler = GradScaler("cuda", enabled=use_amp) | |
| # Resume handling | |
| best_ckpt_path = os.path.join(ckpt_dir, "best_model.pt") | |
| latest_ckpt_path = os.path.join(ckpt_dir, "latest_checkpoint.pt") | |
| start_epoch = 1 | |
| best_metric_val = -1.0 | |
| metrics_csv_path = os.path.join(log_dir, "metrics.csv") | |
| if resume and os.path.exists(latest_ckpt_path): | |
| logger.info(f"Resuming training from checkpoint: {latest_ckpt_path}") | |
| checkpoint = torch.load(latest_ckpt_path, map_location=device) | |
| model.load_state_dict(checkpoint["model_state_dict"]) | |
| optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) | |
| if "scheduler_state_dict" in checkpoint and checkpoint["scheduler_state_dict"]: | |
| scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) | |
| if "scaler_state_dict" in checkpoint and checkpoint["scaler_state_dict"]: | |
| scaler.load_state_dict(checkpoint["scaler_state_dict"]) | |
| start_epoch = checkpoint["epoch"] + 1 | |
| best_metric_val = checkpoint.get("best_metric_val", -1.0) | |
| logger.info(f"Resumed at epoch {start_epoch} (Previous Best Metric: {best_metric_val:.4f})") | |
| else: | |
| # Initialize metrics CSV | |
| if not os.path.exists(metrics_csv_path): | |
| pd.DataFrame(columns=[ | |
| "epoch", "train_loss", "val_loss", | |
| "val_precision", "val_sensitivity", "val_specificity", "val_f1", | |
| "val_pr_auc", "val_roc_auc", "epoch_time_sec", "is_best" | |
| ]).to_csv(metrics_csv_path, index=False) | |
| # Early stopping settings | |
| patience = train_cfg.get("early_stopping_patience", 5) | |
| selection_metric = train_cfg.get("early_stopping_metric", "val_pr_auc") | |
| no_improve_epochs = 0 | |
| grad_clip = train_cfg.get("grad_clip_max_norm", 1.0) | |
| total_train_start = time.time() | |
| # Training Loop | |
| logger.info(f"Starting training loop: Epochs {start_epoch} to {num_epochs}") | |
| for epoch in range(start_epoch, num_epochs + 1): | |
| epoch_start = time.time() | |
| model.train() | |
| running_loss = 0.0 | |
| batch_count = 0 | |
| for batch_idx, (x, y, _) in enumerate(train_loader): | |
| x = x.to(device, non_blocking=True) | |
| y = y.to(device, non_blocking=True) | |
| optimizer.zero_grad() | |
| with autocast(device_type=device.type, enabled=use_amp): | |
| logits = model(x) | |
| loss = criterion(logits, y) | |
| scaler.scale(loss).backward() | |
| if grad_clip: | |
| scaler.unscale_(optimizer) | |
| nn.utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip) | |
| scaler.step(optimizer) | |
| scaler.update() | |
| running_loss += loss.item() * len(y) | |
| batch_count += 1 | |
| if (batch_idx + 1) % 50 == 0 or (batch_idx + 1) == len(train_loader): | |
| logger.info(f" Epoch {epoch}/{num_epochs} [{batch_idx+1}/{len(train_loader)}] Step Loss: {loss.item():.4f}") | |
| scheduler.step() | |
| train_loss = running_loss / len(train_loader.dataset) | |
| # Validation phase (strictly float32 precision, matching evaluate.py) | |
| val_metrics = evaluate_split(model, val_loader, criterion, device, threshold=0.5) | |
| epoch_time = time.time() - epoch_start | |
| current_val_metric = val_metrics[selection_metric.replace("val_", "")] | |
| is_best = current_val_metric > best_metric_val | |
| logger.info( | |
| f"Epoch {epoch:02d}/{num_epochs:02d} Complete ({epoch_time:.1f}s) | " | |
| f"Train Loss: {train_loss:.4f} | Val Loss: {val_metrics['loss']:.4f} | " | |
| f"PR-AUC: {val_metrics['pr_auc']:.4f} | F1: {val_metrics['f1']:.4f} | " | |
| f"Sens: {val_metrics['sensitivity']:.4f} | Spec: {val_metrics['specificity']:.4f}" | |
| ) | |
| # Save Latest Checkpoint | |
| latest_checkpoint_data = { | |
| "epoch": epoch, | |
| "model_state_dict": model.state_dict(), | |
| "optimizer_state_dict": optimizer.state_dict(), | |
| "scheduler_state_dict": scheduler.state_dict(), | |
| "scaler_state_dict": scaler.state_dict(), | |
| "best_metric_val": best_metric_val if not is_best else current_val_metric, | |
| "config": config, | |
| "val_metrics": val_metrics, | |
| "train_loss": train_loss, | |
| "timestamp": datetime.utcnow().isoformat() | |
| } | |
| torch.save(latest_checkpoint_data, latest_ckpt_path) | |
| # Save Best Checkpoint | |
| if is_best: | |
| best_metric_val = current_val_metric | |
| no_improve_epochs = 0 | |
| torch.save(latest_checkpoint_data, best_ckpt_path) | |
| logger.info(f" >>> NEW BEST MODEL SAVED ({selection_metric}: {best_metric_val:.4f}) to {best_ckpt_path}") | |
| else: | |
| no_improve_epochs += 1 | |
| logger.info(f" No improvement in {selection_metric} for {no_improve_epochs} epoch(s). Best: {best_metric_val:.4f}") | |
| # Append to CSV | |
| metric_row = pd.DataFrame([{ | |
| "epoch": epoch, | |
| "train_loss": round(train_loss, 4), | |
| "val_loss": round(val_metrics["loss"], 4), | |
| "val_precision": round(val_metrics["precision"], 4), | |
| "val_sensitivity": round(val_metrics["sensitivity"], 4), | |
| "val_specificity": round(val_metrics["specificity"], 4), | |
| "val_f1": round(val_metrics["f1"], 4), | |
| "val_pr_auc": round(val_metrics["pr_auc"], 4), | |
| "val_roc_auc": round(val_metrics["roc_auc"], 4), | |
| "epoch_time_sec": round(epoch_time, 2), | |
| "is_best": is_best | |
| }]) | |
| metric_row.to_csv(metrics_csv_path, mode="a", header=False, index=False) | |
| # Early stopping check | |
| if no_improve_epochs >= patience: | |
| logger.info(f"Early stopping triggered after {patience} epochs without improvement in {selection_metric}!") | |
| break | |
| total_duration = time.time() - total_train_start | |
| logger.info("=" * 70) | |
| logger.info(f"TRAINING COMPLETE in {total_duration:.1f}s ({total_duration/60:.2f} mins)") | |
| logger.info(f"Best {selection_metric}: {best_metric_val:.4f}") | |
| logger.info(f"Best checkpoint: {best_ckpt_path}") | |
| logger.info(f"Metrics CSV: {metrics_csv_path}") | |
| return { | |
| "best_ckpt_path": best_ckpt_path, | |
| "best_metric_val": best_metric_val, | |
| "total_duration": total_duration, | |
| "epochs_completed": epoch | |
| } | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description="Train baseline CNN-LSTM for EEG Seizure Detection") | |
| parser.add_argument("--config", default="configs/baseline_config.json", help="Path to config JSON") | |
| parser.add_argument("--resume", action="store_true", help="Resume training from latest checkpoint") | |
| args = parser.parse_args() | |
| train_pipeline(config_path=args.config, resume=args.resume) | |
Xet Storage Details
- Size:
- 15.3 kB
- Xet hash:
- 14450faf05a4ed1ae3c75e6f56b6aaef40080f79c6ed5104fd76e4af2887e111
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.