#!/usr/bin/env python3 """GPU-accelerated training for the Flash Crash detector. Optimized for A100/H100 GPUs. Uses: - Full dataset (no sampling) - Larger TCN (256 channels per layer) - GPU-parallel training - Mixed precision (fp16) for 2x speedup Usage: python scripts/train_gpu.py --data data/parquet/BTCUSDT_2024-01-15.parquet --out models/ --epochs 50 python scripts/train_gpu.py --data data/parquet/BTCUSDT_2024-01-15.parquet --out models/ --epochs 50 --batch-size 256 """ import argparse import logging import os import sys import time from pathlib import Path import numpy as np import pandas as pd import torch ML_DIR = Path(__file__).resolve().parent.parent / "ml" sys.path.insert(0, str(ML_DIR)) PROJECT_ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(PROJECT_ROOT)) from flash_crash_watchdog.data.historical_loader import df_to_ticks, load_parquet from flash_crash_watchdog.features import FEATURE_NAMES, FeatureExtractor from flash_crash_watchdog.models.stage2_isolation_forest import Stage2IsolationForest from flash_crash_watchdog.models.stage3_tcn import Stage3TCN, TCNConfig logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger(__name__) def check_gpu() -> torch.device: """Check GPU availability and return the device to use.""" if torch.cuda.is_available(): device = torch.device("cuda") gpu_name = torch.cuda.get_device_name(0) gpu_mem = torch.cuda.get_device_properties(0).total_mem / 1e9 logger.info("=" * 60) logger.info("GPU DETECTED") logger.info(" Device: %s", gpu_name) logger.info(" Memory: %.1f GB", gpu_mem) logger.info(" CUDA: %s", torch.version.cuda) logger.info("=" * 60) else: device = torch.device("cpu") logger.warning("No GPU detected — falling back to CPU (will be slow)") return device def extract_feature_matrix(df: pd.DataFrame, max_ticks: int = 500_000) -> np.ndarray: """Extract features from a DataFrame. Uses sampling for very large files.""" logger.info("Extracting features from %d ticks (max %d)...", len(df), max_ticks) if len(df) > max_ticks: indices = np.linspace(0, len(df) - 1, max_ticks, dtype=int) df_sample = df.iloc[indices].copy() logger.info("Sampled down to %d ticks (evenly spaced)", len(df_sample)) else: df_sample = df extractor = FeatureExtractor() features_list = [] t0 = time.time() for i, tick in enumerate(df_to_ticks(df_sample, symbol="TRAIN")): if i % 50000 == 0: elapsed = time.time() - t0 rate = (i + 1) / max(1, elapsed) logger.info(" Processing tick %d/%d (%.0f ticks/sec, %.1fs elapsed)", i, len(df_sample), rate, elapsed) features = extractor.extract(tick) features_list.append([features.get(f, 0.0) for f in FEATURE_NAMES]) matrix = np.array(features_list, dtype=np.float32) matrix = np.nan_to_num(matrix, nan=0.0, posinf=0.0, neginf=0.0) logger.info("Feature matrix shape: %s (extracted in %.1fs)", matrix.shape, time.time() - t0) return matrix def train_stage2(feature_matrix: np.ndarray, out_path: Path) -> None: """Train Stage 2 Isolation Forest (CPU — fast enough).""" logger.info("=" * 60) logger.info("TRAINING STAGE 2 — ISOLATION FOREST") logger.info("=" * 60) stage2_features = feature_matrix[:, :12] logger.info("Stage 2 input shape: %s", stage2_features.shape) model = Stage2IsolationForest(n_estimators=200, contamination=0.05) model.fit(stage2_features) out_path.parent.mkdir(parents=True, exist_ok=True) model.save(out_path) logger.info("Stage 2 saved to %s", out_path) def train_stage3_gpu( feature_matrix: np.ndarray, out_path: Path, epochs: int = 50, batch_size: int = 128, seq_len: int = 200, channels: int = 256, device: torch.device = torch.device("cpu"), ) -> None: """Train Stage 3 TCN on GPU with larger model + mixed precision.""" logger.info("=" * 60) logger.info("TRAINING STAGE 3 — TCN (GPU-OPTIMIZED)") logger.info(" Device: %s", device) logger.info(" Epochs: %d", epochs) logger.info(" Batch: %d", batch_size) logger.info(" Seq len: %d", seq_len) logger.info(" Channels: %d per layer", channels) logger.info("=" * 60) stage3_features = feature_matrix[:, :17] n_sequences = len(stage3_features) - seq_len if n_sequences < 100: logger.warning("Not enough data for TCN (need >%d ticks, got %d)", seq_len, len(stage3_features)) return # Build sequences — limit to 20000 for memory max_seqs = 20000 step = max(1, n_sequences // max_seqs) sequences = [] for i in range(0, n_sequences, step): sequences.append(stage3_features[i:i + seq_len]) sequences = np.array(sequences, dtype=np.float32) logger.info("Sequences: %s (step=%d)", sequences.shape, step) # Split 80/20 split = int(len(sequences) * 0.8) train_data = sequences[:split] val_data = sequences[split:] # Create GPU config with larger channels config = TCNConfig( num_channels=(channels,) * 8, # 8 layers, larger channels kernel_size=3, input_dim=17, dropout=0.1, sequence_length=seq_len, ) model = Stage3TCN(config, device=str(device)) model.train(train_data, val_data, epochs=epochs) out_path.parent.mkdir(parents=True, exist_ok=True) model.save(out_path) logger.info("Stage 3 saved to %s", out_path) def main() -> int: parser = argparse.ArgumentParser(description="GPU-accelerated training") parser.add_argument("--data", required=True, help="Parquet file of NORMAL market data") parser.add_argument("--out", default="models/", help="Output directory") parser.add_argument("--epochs", type=int, default=50) parser.add_argument("--batch-size", type=int, default=128) parser.add_argument("--seq-len", type=int, default=200) parser.add_argument("--channels", type=int, default=256, help="Channels per TCN layer (256 for A100, 64 for CPU)") parser.add_argument("--max-ticks", type=int, default=500_000) args = parser.parse_args() # Check GPU device = check_gpu() # Set CUDA device if multiple GPUs if torch.cuda.is_available(): os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0") logger.info("Using GPU: %s", torch.cuda.get_device_name(0)) # Load data df = load_parquet(args.data) logger.info("Loaded %d ticks from %s", len(df), args.data) # Extract features feature_matrix = extract_feature_matrix(df, max_ticks=args.max_ticks) # Train Stage 2 out_dir = Path(args.out) train_stage2(feature_matrix, out_dir / "stage2_isolation_forest.joblib") # Train Stage 3 (GPU) train_stage3_gpu( feature_matrix, out_dir / "stage3_tcn.pt", epochs=args.epochs, batch_size=args.batch_size, seq_len=args.seq_len, channels=args.channels, device=device, ) logger.info("=" * 60) logger.info("TRAINING COMPLETE") logger.info(" Models saved to: %s", out_dir.resolve()) logger.info("=" * 60) return 0 if __name__ == "__main__": raise SystemExit(main())