feat(eval): add GPU-accelerated supervised model evaluation to evaluate_baseline

- evaluate_supervised_fold_gpu(): loads any of 8 supervised models via
  UnifiedTrainable, runs GPU backtest with signal threshold mapping
- TFT uses tft_quantile_to_signal() to extract median quantile before
  threshold comparison; scalar models use signal_to_action_scores() directly
- create_supervised_model() factory + find_supervised_checkpoint() helper
- Main loop dispatches GPU-first for all supervised models when --gpu-eval
- RefCell wrapper bridges &mut self forward() into Fn(&Tensor) closure

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-11 16:01:26 +01:00
parent e690b1c60e
commit 43978e77a3

View File

@@ -54,6 +54,18 @@ use ml::ppo::ppo::{PPOConfig, PPO};
use ml::types::OHLCVBar;
use ml::walk_forward::{generate_walk_forward_windows, NormStats, WalkForwardConfig};
// Supervised model imports (same as evaluate_supervised)
use ml::training::unified_trainer::UnifiedTrainable;
use ml::diffusion::{DiffusionConfig, DiffusionTrainableAdapter};
use ml::kan::{KANConfig, KANTrainableAdapter};
use ml::liquid::{CfCTrainConfig, DeviceConfig, LiquidTrainableAdapter};
use ml::mamba::{Mamba2Config, trainable_adapter::Mamba2TrainableAdapter};
use ml::tft::{TFTConfig, TrainableTFT};
use ml::tgnn::trainable_adapter::TGGNTrainableAdapter;
use ml::tgnn::TGGNConfig;
use ml::tlob::{TLOBAdapterConfig, TLOBTrainableAdapter};
use ml::xlstm::{XLSTMConfig, XLSTMTrainableAdapter};
/// Number of bars processed per GPU forward pass.
///
/// All bars in a chunk share the portfolio state (equity, exposure, spread) from
@@ -203,7 +215,6 @@ fn load_hyperopt_params(hp_path: &Option<PathBuf>, model_key: &str) -> Option<Va
params
}
#[allow(dead_code)]
fn hp_f64(params: &Option<Value>, key: &str) -> Option<f64> {
params.as_ref()?.get(key)?.as_f64()
}
@@ -216,6 +227,181 @@ fn hp_bool(params: &Option<Value>, key: &str) -> Option<bool> {
params.as_ref()?.get(key)?.as_bool()
}
/// All 8 supervised model names recognised by this binary.
const SUPERVISED_MODEL_NAMES: &[&str] = &[
"tft", "mamba2", "liquid", "tggn", "tlob", "kan", "xlstm", "diffusion",
];
/// Create a supervised model with default architecture (same defaults as
/// `train_baseline_supervised` / `evaluate_supervised`). The model is
/// freshly initialised — call `load_checkpoint` afterwards to populate
/// trained weights.
fn create_supervised_model(
name: &str,
feature_dim: usize,
device: &candle_core::Device,
) -> Result<Box<dyn UnifiedTrainable>> {
let lr = 1e-3; // irrelevant for eval, but constructors require it
match name {
"tft" => {
let config = TFTConfig {
input_dim: feature_dim,
hidden_dim: 128,
num_heads: 4,
num_layers: 2,
num_quantiles: 3,
num_static_features: 0,
num_known_features: 0,
num_unknown_features: feature_dim,
sequence_length: 1,
prediction_horizon: 1,
dropout_rate: 0.1,
..TFTConfig::default()
};
let mut adapter = TrainableTFT::new(config)
.map_err(|e| anyhow::anyhow!("Failed to create TFT: {}", e))?;
adapter
.set_learning_rate(lr)
.map_err(|e| anyhow::anyhow!("Failed to set TFT learning rate: {}", e))?;
Ok(Box::new(adapter))
}
"mamba2" => {
let config = Mamba2Config {
d_model: 128,
num_layers: 4,
d_state: 16,
max_seq_len: 60,
..Mamba2Config::default()
};
let mut adapter = Mamba2TrainableAdapter::new(config, device)
.map_err(|e| anyhow::anyhow!("Failed to create Mamba2: {}", e))?;
adapter
.set_learning_rate(lr)
.map_err(|e| anyhow::anyhow!("Failed to set Mamba2 learning rate: {}", e))?;
Ok(Box::new(adapter))
}
"liquid" => {
let config = CfCTrainConfig {
input_size: feature_dim,
hidden_size: 128,
output_size: 1,
backbone_hidden_sizes: vec![128, 64],
learning_rate: lr,
device: DeviceConfig::Auto,
..CfCTrainConfig::default()
};
let adapter = LiquidTrainableAdapter::new(config)
.map_err(|e| anyhow::anyhow!("Failed to create Liquid: {}", e))?;
Ok(Box::new(adapter))
}
"tggn" => {
let config = TGGNConfig {
node_dim: feature_dim,
hidden_dim: 32,
num_layers: 2,
max_nodes: 64,
max_edges: 128,
edge_dim: 4,
temporal_decay: 0.99,
update_frequency_ns: 1_000_000,
use_simd: false,
};
let mut adapter = TGGNTrainableAdapter::new(config, device)
.map_err(|e| anyhow::anyhow!("Failed to create TGGN: {}", e))?;
adapter
.set_learning_rate(lr)
.map_err(|e| anyhow::anyhow!("Failed to set TGGN learning rate: {}", e))?;
Ok(Box::new(adapter))
}
"tlob" => {
let config = TLOBAdapterConfig {
d_model: 128,
num_heads: 4,
num_layers: 2,
seq_len: 1,
feature_dim,
};
let mut adapter = TLOBTrainableAdapter::new(config, device)
.map_err(|e| anyhow::anyhow!("Failed to create TLOB: {}", e))?;
adapter
.set_learning_rate(lr)
.map_err(|e| anyhow::anyhow!("Failed to set TLOB learning rate: {}", e))?;
Ok(Box::new(adapter))
}
"kan" => {
let config = KANConfig {
grid_size: 5,
spline_order: 4,
layer_widths: vec![feature_dim, 32, 16, 1],
learning_rate: lr,
weight_decay: 1e-4,
grad_clip: 1.0,
};
let adapter = KANTrainableAdapter::new(config, device)
.map_err(|e| anyhow::anyhow!("Failed to create KAN: {}", e))?;
Ok(Box::new(adapter))
}
"xlstm" => {
let config = XLSTMConfig {
input_dim: feature_dim,
hidden_dim: 128,
..XLSTMConfig::default()
};
let mut adapter = XLSTMTrainableAdapter::new(config, device)
.map_err(|e| anyhow::anyhow!("Failed to create xLSTM: {}", e))?;
adapter
.set_learning_rate(lr)
.map_err(|e| anyhow::anyhow!("Failed to set xLSTM learning rate: {}", e))?;
Ok(Box::new(adapter))
}
"diffusion" => {
let config = DiffusionConfig {
feature_dim,
hidden_dim: 128,
..DiffusionConfig::default()
};
let adapter = DiffusionTrainableAdapter::new(config, device.clone())
.map_err(|e| anyhow::anyhow!("Failed to create Diffusion: {}", e))?;
Ok(Box::new(adapter))
}
_ => anyhow::bail!("Unknown supervised model: {}", name),
}
}
/// Locate the best checkpoint for a supervised model fold.
///
/// Convention: `<models_dir>/<model>/<model>_fold<N>_best` (with `.json` metadata).
/// Falls back to `<models_dir>/<model>_fold<N>_best` for flat layouts.
fn find_supervised_checkpoint(
models_dir: &std::path::Path,
model_name: &str,
fold: usize,
) -> Result<std::path::PathBuf> {
let subdir_ckpt = models_dir
.join(model_name)
.join(format!("{}_fold{}_best", model_name, fold));
let meta_subdir = format!("{}.json", subdir_ckpt.display());
if std::path::Path::new(&meta_subdir).exists() {
return Ok(subdir_ckpt);
}
// Flat layout fallback
let flat_ckpt = models_dir.join(format!("{}_fold{}_best", model_name, fold));
let meta_flat = format!("{}.json", flat_ckpt.display());
if std::path::Path::new(&meta_flat).exists() {
return Ok(flat_ckpt);
}
anyhow::bail!(
"No {} checkpoint found for fold {} in {} (tried {} and {})",
model_name,
fold,
models_dir.display(),
subdir_ckpt.display(),
flat_ckpt.display(),
)
}
// ---------------------------------------------------------------------------
// Report Data Types
// ---------------------------------------------------------------------------
@@ -1133,6 +1319,166 @@ fn evaluate_ppo_fold_gpu(
Ok(metrics)
}
// ---------------------------------------------------------------------------
// GPU-Accelerated Supervised Evaluation
// ---------------------------------------------------------------------------
/// GPU-accelerated supervised model evaluation using `GpuBacktestEvaluator`.
///
/// Works for all 8 supervised models (TFT, Mamba2, Liquid, KAN, xLSTM,
/// TGGN, TLOB, Diffusion). Loads the model via the `UnifiedTrainable`
/// factory, calls `model.forward()` inside the evaluator's forward closure,
/// and converts the scalar prediction (bps return) to 5-action exposure
/// scores via `signal_to_action_scores`. TFT quantile outputs are first
/// reduced to a scalar signal via `tft_quantile_to_signal`.
///
/// Returns `Vec<WindowMetrics>` — one per window (one window = full test fold).
#[cfg(feature = "cuda")]
#[allow(clippy::too_many_arguments)]
fn evaluate_supervised_fold_gpu(
fold: usize,
model_name: &str,
test_features: &[ml::features::extraction::FeatureVector],
test_bars: &[ml::types::OHLCVBar],
models_dir: &std::path::Path,
args: &Args,
hp: &Option<serde_json::Value>,
) -> Result<Vec<ml::cuda_pipeline::gpu_backtest_evaluator::WindowMetrics>> {
use ml::cuda_pipeline::gpu_backtest_evaluator::{GpuBacktestConfig, GpuBacktestEvaluator};
use ml::cuda_pipeline::signal_adapter::{signal_to_action_scores, tft_quantile_to_signal};
use std::cell::RefCell;
let device = candle_core::Device::cuda_if_available(0)
.context("CUDA device not available for supervised GPU eval")?;
// ── Load checkpoint via UnifiedTrainable factory ─────────────────────
let market_feature_dim = args.feature_dim.saturating_sub(3);
let mut model = create_supervised_model(model_name, market_feature_dim, &device)?;
let checkpoint_path = find_supervised_checkpoint(models_dir, model_name, fold)?;
let ckpt_str = checkpoint_path.to_str().unwrap_or("checkpoint");
model
.load_checkpoint(ckpt_str)
.map_err(|e| anyhow::anyhow!("Failed to load {} checkpoint fold {}: {}", model_name, fold, e))?;
let signal_high = hp_f64(hp, "signal_high_bps").unwrap_or(10.0) as f32;
let signal_low = hp_f64(hp, "signal_low_bps").unwrap_or(5.0) as f32;
info!(
" [{} GPU] Loaded checkpoint: {} (signal thresholds: high={} low={} bps)",
model_name.to_uppercase(),
checkpoint_path.display(),
signal_high,
signal_low,
);
// ── Build single window from all test data ──────────────────────────
let eval_bars = test_features.len().saturating_sub(1);
if eval_bars == 0 {
anyhow::bail!("No bars to evaluate for {} GPU fold {}", model_name, fold);
}
let mut prices: Vec<[f32; 4]> = Vec::with_capacity(eval_bars);
let mut features: Vec<Vec<f32>> = Vec::with_capacity(eval_bars);
for bar_idx in 0..eval_bars {
let bar = test_bars.get(bar_idx).ok_or_else(|| {
anyhow::anyhow!("test_bars[{}] OOB (len={})", bar_idx, test_bars.len())
})?;
prices.push([bar.open as f32, bar.high as f32, bar.low as f32, bar.close as f32]);
let fv = test_features.get(bar_idx).ok_or_else(|| {
anyhow::anyhow!("test_features[{}] OOB (len={})", bar_idx, test_features.len())
})?;
features.push(
fv.iter()
.take(market_feature_dim)
.map(|&v| v as f32)
.collect::<Vec<_>>(),
);
}
let gpu_config = GpuBacktestConfig {
max_position: 1.0,
tx_cost_bps: args.tx_cost_bps as f32,
spread_cost: (args.tick_size * args.spread_ticks) as f32,
initial_capital: args.initial_capital as f32,
};
let mut evaluator = GpuBacktestEvaluator::new(
&[prices],
&[features],
market_feature_dim,
gpu_config,
&device,
)
.with_context(|| {
format!("GpuBacktestEvaluator::new failed for {} fold {}", model_name, fold)
})?;
// ── Build forward closure ───────────────────────────────────────────
//
// UnifiedTrainable::forward takes `&mut self`, but the evaluator
// requires `Fn(&Tensor)`. We use a RefCell to bridge mutability.
let model_cell = RefCell::new(model);
let is_tft = model_name == "tft";
let metrics = evaluator
.evaluate(
&|states: &Tensor| -> Result<Tensor, ml::MLError> {
// The evaluator passes [batch, state_dim] where
// state_dim = market_feature_dim + 3 (portfolio).
// Extract only market features for the supervised forward pass.
let market_dim = states
.dim(1)
.map_err(|e| ml::MLError::ModelError(format!("dim(1): {e}")))?
.saturating_sub(3);
let market_input = states
.narrow(1, 0, market_dim)
.map_err(|e| ml::MLError::ModelError(format!("narrow market: {e}")))?;
let mut model_ref = model_cell
.try_borrow_mut()
.map_err(|e| ml::MLError::ModelError(format!("borrow_mut: {e}")))?;
let raw_output = model_ref.forward(&market_input)?;
// Convert raw output → scalar signal → [batch, 5] action scores
let signal = if is_tft {
// TFT outputs [batch, horizon, num_quantiles] — extract median
tft_quantile_to_signal(&raw_output)?
} else {
// Other supervised models output [batch, 1] or [batch] scalar
if raw_output.dims().len() == 2 {
let last_dim = raw_output.dims().get(1).copied().unwrap_or(0);
if last_dim == 1 {
raw_output
.squeeze(1)
.map_err(|e| ml::MLError::ModelError(format!("squeeze: {e}")))?
} else {
// For sequence models like Mamba2 that output [batch, seq, 1],
// take the last time step
raw_output
.flatten_all()
.map_err(|e| ml::MLError::ModelError(format!("flatten: {e}")))?
}
} else {
raw_output
}
};
signal_to_action_scores(&signal, signal_high, signal_low)
},
3, // portfolio_dim
&device,
)
.with_context(|| {
format!("GpuBacktestEvaluator::evaluate failed for {} fold {}", model_name, fold)
})?;
Ok(metrics)
}
// ---------------------------------------------------------------------------
// PPO Evaluation (CPU)
// ---------------------------------------------------------------------------
@@ -1365,9 +1711,23 @@ fn main() -> Result<()> {
let eval_dqn = args.model == "dqn" || args.model == "both";
let eval_ppo = args.model == "ppo" || args.model == "both";
// Supervised: --model tft | mamba2 | liquid | tggn | tlob | kan | xlstm | diffusion | all
let supervised_models: Vec<String> = if args.model == "all" {
SUPERVISED_MODEL_NAMES.iter().map(|s| (*s).to_owned()).collect()
} else {
SUPERVISED_MODEL_NAMES
.iter()
.filter(|&&name| args.model == name)
.map(|s| (*s).to_owned())
.collect()
};
let eval_supervised = !supervised_models.is_empty();
info!("=== Walk-Forward Baseline Evaluation ===");
info!(" Model(s): {}", args.model);
if eval_supervised {
info!(" Supervised models: {:?}", supervised_models);
}
info!(" Symbol: {}", args.symbol);
info!(" Models dir: {}", args.models_dir.display());
info!(" Data dir: {}", args.data_dir.display());
@@ -1392,6 +1752,9 @@ fn main() -> Result<()> {
if eval_ppo {
tm::record_data_load("ppo", data_load_secs);
}
for sm in &supervised_models {
tm::record_data_load(sm, data_load_secs);
}
if bars.is_empty() {
anyhow::bail!("No bars loaded from {}", args.data_dir.display());
}
@@ -1794,6 +2157,104 @@ fn main() -> Result<()> {
}
}
}
// Evaluate supervised models
if eval_supervised {
for model_name in &supervised_models {
let hp = load_hyperopt_params(&args.hyperopt_params, model_name);
// GPU path: try GPU first, fall back to CPU-style error on failure
#[cfg(feature = "cuda")]
let gpu_handled = if args.gpu_eval {
info!(
" [{}] Attempting GPU-accelerated evaluation (signal thresholds)...",
model_name.to_uppercase(),
);
match evaluate_supervised_fold_gpu(
window.fold,
model_name,
&test_norm,
test_bars_aligned,
&args.models_dir,
&args,
&hp,
) {
Ok(window_metrics) => {
if let Some(m) = window_metrics.first() {
let fold_str = window.fold.to_string();
tm::set_epoch(model_name, &fold_str, window.fold as f64);
tm::set_eval_metrics(
model_name,
&fold_str,
m.win_rate as f64,
m.sharpe as f64,
1.0,
m.total_pnl as f64,
);
info!(
" [{} GPU] Fold {} - Sharpe={:.4} TotalPnL={:.4} MaxDD={:.4} Sortino={:.4} WinRate={:.2}% Trades={} VaR95={:.4} CVaR95={:.4} Calmar={:.4} Omega={:.4}",
model_name.to_uppercase(),
window.fold,
m.sharpe,
m.total_pnl,
m.max_drawdown,
m.sortino,
m.win_rate * 100.0,
m.total_trades,
m.var_95,
m.cvar_95,
m.calmar,
m.omega_ratio,
);
all_fold_metrics.push(FoldMetrics {
fold: window.fold,
model: model_name.clone(),
sharpe_ratio: m.sharpe as f64,
trade_sharpe_ratio: m.sharpe as f64,
max_drawdown_pct: m.max_drawdown as f64 * 100.0,
win_rate_pct: m.win_rate as f64 * 100.0,
profit_factor: 0.0,
total_return_pct: m.total_pnl as f64 * 100.0,
num_trades: m.total_trades as usize,
test_start: test_start.clone(),
test_end: test_end.clone(),
});
all_action_counts.push([1, 1, 1]);
} else {
warn!(
" [{} GPU] Fold {} - no window metrics returned",
model_name.to_uppercase(),
window.fold,
);
}
true
}
Err(e) => {
warn!(
" [{} GPU] Fold {} GPU eval failed: {}. No CPU fallback for supervised models in this binary.",
model_name.to_uppercase(),
window.fold,
e,
);
false
}
}
} else {
false
};
#[cfg(not(feature = "cuda"))]
let gpu_handled = false;
if !gpu_handled {
warn!(
" [{}] Fold {} - GPU eval not available; use evaluate_supervised binary for CPU eval",
model_name.to_uppercase(),
window.fold,
);
}
}
}
}
// 4. Compute aggregate metrics