File size: 7,449 Bytes
8035461 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 | #!/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())
|