bcf9ecd07c
TFT training completeness: - Fix mod.rs::train() stub backward → real AdamW optimizer + gradient flow - Fix TFTTrainer optimizer init, backward pass, LR scheduling, checkpointing - Fix temporal_attention weight tracking (RwLock, stores per-head means) Mamba2 discretization: - Replace raw continuous-time state transition with proper discretization - SSD layer now uses softplus(delta) step size with 2nd-order Taylor approximation of matrix exponential: A_bar ≈ I + A*dt + (A*dt)²/2 - Correct dtype handling (F64 SSM matrices, F32 output) PPO entropy fix: - Fix LSTM training path: was using constant entropy (coeff * 0.5), now computes real entropy from log-probabilities Circuit breaker consolidation: - Move canonical implementation to ml/src/common/circuit_breaker.rs - DQN and PPO circuit_breaker.rs now re-export from common Validation stack additions: - Add CPCV (Combinatorial Purged Cross-Validation) with purging/embargo - Add FDR correction (Benjamini-Hochberg + Benjamini-Yekutieli) 1922 lib tests pass, 0 failures. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1510 lines
52 KiB
Rust
1510 lines
52 KiB
Rust
//! # Temporal Fusion Transformer (TFT) for HFT
|
|
//!
|
|
//! State-of-the-art multi-horizon forecasting with variable selection networks,
|
|
//! temporal self-attention, gated residual networks, and uncertainty quantification.
|
|
//!
|
|
//! ## Key Features
|
|
//!
|
|
//! - Multi-horizon forecasting (1-tick to 100-tick ahead)
|
|
//! - Variable selection networks for feature importance
|
|
//! - Gated residual networks for improved gradient flow
|
|
//! - Quantile outputs for uncertainty estimation
|
|
//! - Temporal self-attention for sequential modeling
|
|
//! - Sub-50μs inference latency optimized for HFT
|
|
//!
|
|
//! ## Performance Targets
|
|
//!
|
|
//! - Inference: <50μs per prediction
|
|
//! - Accuracy improvement: +15% over baseline
|
|
//! - Memory usage: <1GB
|
|
//! - Throughput: >100K predictions/sec
|
|
|
|
use std::collections::HashMap;
|
|
use std::num::NonZeroUsize;
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
use std::sync::Arc;
|
|
use std::time::{Instant, SystemTime};
|
|
|
|
use async_trait::async_trait;
|
|
use candle_core::{DType, Device, Module, Tensor};
|
|
use candle_nn::{linear, AdamW, Linear, Optimizer, ParamsAdamW, VarBuilder, VarMap};
|
|
use lru::LruCache;
|
|
use ndarray::{Array1, Array2};
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::Value;
|
|
use tracing::{debug, info, instrument, warn};
|
|
use uuid::Uuid;
|
|
|
|
use crate::checkpoint::Checkpointable;
|
|
use crate::{MLError, ModelType};
|
|
|
|
// Import TFT components
|
|
pub mod gated_residual;
|
|
pub mod hft_optimizations;
|
|
pub mod lstm_encoder;
|
|
pub mod qat_tft; // Quantization-Aware Training wrapper - RE-ENABLED: Device mismatch fix applied
|
|
pub mod quantile_outputs;
|
|
pub mod quantized_attention; // Re-enabled Wave 9.12
|
|
pub mod quantized_grn;
|
|
pub mod quantized_lstm;
|
|
pub mod quantized_tft; // Re-enabled Wave 9.12
|
|
pub mod quantized_vsn;
|
|
pub mod temporal_attention;
|
|
pub mod trainable_adapter;
|
|
pub mod training;
|
|
pub mod variable_selection;
|
|
pub mod varmap_quantization;
|
|
|
|
// Public exports for TFT components
|
|
pub use gated_residual::{GRNStack, GatedResidualNetwork};
|
|
pub use lstm_encoder::LSTMEncoder;
|
|
pub use qat_tft::QATTemporalFusionTransformer; // Quantization-Aware Training wrapper - RE-ENABLED: Device mismatch fix applied
|
|
pub use quantile_outputs::QuantileLayer;
|
|
pub use quantized_attention::QuantizedTemporalAttention; // Re-enabled Wave 9.12
|
|
pub use quantized_grn::QuantizedGatedResidualNetwork;
|
|
pub use quantized_lstm::QuantizedLSTMEncoder;
|
|
pub use quantized_tft::QuantizedTemporalFusionTransformer; // Re-enabled Wave 9.12
|
|
pub use quantized_vsn::QuantizedVariableSelectionNetwork;
|
|
pub use temporal_attention::TemporalSelfAttention;
|
|
pub use trainable_adapter::TrainableTFT;
|
|
pub use variable_selection::VariableSelectionNetwork;
|
|
pub use varmap_quantization::{
|
|
load_quantized_weights, quantize_varmap, quantize_varmap_parallel, save_quantized_weights,
|
|
};
|
|
|
|
/// `TFT` Configuration
|
|
/// TFT model variant selection (F32 vs INT8)
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum TFTVariant {
|
|
/// Full precision (F32) model
|
|
F32,
|
|
/// INT8 quantized model (75% memory reduction)
|
|
INT8,
|
|
}
|
|
|
|
impl Default for TFTVariant {
|
|
fn default() -> Self {
|
|
Self::F32
|
|
}
|
|
}
|
|
|
|
impl TFTVariant {
|
|
/// Check if variant uses quantization
|
|
pub fn is_quantized(&self) -> bool {
|
|
matches!(self, Self::INT8)
|
|
}
|
|
|
|
/// Get expected memory reduction ratio vs F32
|
|
pub fn memory_reduction_ratio(&self) -> f64 {
|
|
match self {
|
|
Self::F32 => 1.0,
|
|
Self::INT8 => 0.25, // 75% reduction → 25% of original
|
|
}
|
|
}
|
|
}
|
|
|
|
/// `TFT` Configuration
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TFTConfig {
|
|
// Model architecture
|
|
pub input_dim: usize,
|
|
pub hidden_dim: usize,
|
|
pub num_heads: usize,
|
|
pub num_layers: usize,
|
|
|
|
// Forecasting parameters
|
|
pub prediction_horizon: usize,
|
|
pub sequence_length: usize,
|
|
pub num_quantiles: usize,
|
|
|
|
// Feature types
|
|
pub num_static_features: usize,
|
|
pub num_known_features: usize,
|
|
pub num_unknown_features: usize,
|
|
|
|
// Training parameters
|
|
pub learning_rate: f64,
|
|
pub batch_size: usize,
|
|
pub dropout_rate: f64,
|
|
pub l2_regularization: f64,
|
|
|
|
// HFT optimization
|
|
pub use_flash_attention: bool,
|
|
pub mixed_precision: bool,
|
|
pub memory_efficient: bool,
|
|
|
|
// Performance constraints
|
|
pub max_inference_latency_us: u64,
|
|
pub target_throughput_pps: u64,
|
|
}
|
|
|
|
impl Default for TFTConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
// Wave C+D: 225 features (201 Wave C + 24 Wave D)
|
|
// Wave C: 201 features (indices 0-200)
|
|
// Wave D: 24 features (indices 201-224)
|
|
input_dim: 225,
|
|
hidden_dim: 128,
|
|
num_heads: 8,
|
|
num_layers: 3,
|
|
prediction_horizon: 10,
|
|
sequence_length: 50,
|
|
num_quantiles: 9,
|
|
// Feature split for 225 total features:
|
|
// - Static: 5 features (symbol metadata)
|
|
// - Known: 10 features (future time features)
|
|
// - Unknown: 210 features (historical OHLCV + technical + microstructure + regime)
|
|
num_static_features: 5,
|
|
num_known_features: 10,
|
|
num_unknown_features: 210,
|
|
learning_rate: 1e-3,
|
|
batch_size: 64,
|
|
dropout_rate: 0.1,
|
|
l2_regularization: 1e-4,
|
|
use_flash_attention: true,
|
|
mixed_precision: true,
|
|
memory_efficient: true,
|
|
max_inference_latency_us: 50,
|
|
target_throughput_pps: 100_000,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// `TFT` Model State for incremental processing
|
|
///
|
|
/// **MEMORY SAFETY FIX (2025-10-25)**:
|
|
/// - Replaced unbounded HashMap with LRU cache (max 1000 entries)
|
|
/// - Prevents 3.6GB/hour memory leak in production inference
|
|
/// - Automatically evicts oldest cache entries when full
|
|
/// - Tested: 1-hour inference run with 10K predictions = stable 24MB memory
|
|
#[derive(Debug, Clone)]
|
|
pub struct TFTState {
|
|
pub hidden_state: Option<Tensor>,
|
|
pub attention_cache: LruCache<String, Tensor>,
|
|
pub last_update: u64,
|
|
}
|
|
|
|
impl TFTState {
|
|
/// Maximum attention cache entries (2000 = ~48MB for TFT-225, 60% training speedup)
|
|
/// Chosen to balance:
|
|
/// - Memory safety: <100MB cache overhead (acceptable for training)
|
|
/// - Hit rate: >95% for typical 50-sequence inference
|
|
/// - Eviction overhead: <0.5% latency impact (reduced by 2x cache size)
|
|
pub const MAX_CACHE_ENTRIES: usize = 2000;
|
|
|
|
pub fn zeros(_config: &TFTConfig) -> Result<Self, MLError> {
|
|
// SAFETY: MAX_CACHE_ENTRIES (2000) is non-zero by construction
|
|
let capacity =
|
|
NonZeroUsize::new(Self::MAX_CACHE_ENTRIES).expect("MAX_CACHE_ENTRIES must be non-zero");
|
|
|
|
Ok(Self {
|
|
hidden_state: None,
|
|
attention_cache: LruCache::new(capacity),
|
|
last_update: 0,
|
|
})
|
|
}
|
|
|
|
/// Clear attention cache to free memory
|
|
/// Call this after training/inference batch to prevent memory accumulation
|
|
pub fn clear_cache(&mut self) {
|
|
self.attention_cache.clear();
|
|
self.hidden_state = None;
|
|
}
|
|
}
|
|
|
|
/// `TFT` Model Metadata
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TFTMetadata {
|
|
pub model_id: String,
|
|
pub version: String,
|
|
pub input_dim: usize,
|
|
pub output_dim: usize,
|
|
pub created_at: SystemTime,
|
|
pub last_trained: Option<SystemTime>,
|
|
pub training_samples: u64,
|
|
pub performance_metrics: HashMap<String, f64>,
|
|
}
|
|
|
|
/// Multi-horizon prediction result
|
|
#[derive(Debug, Clone)]
|
|
pub struct MultiHorizonPrediction {
|
|
pub predictions: Vec<f64>, // Point predictions for each horizon
|
|
pub quantiles: Vec<Vec<f64>>, // Quantile predictions [horizon][quantile]
|
|
pub uncertainty: Vec<f64>, // Uncertainty estimates
|
|
pub confidence_intervals: Vec<(f64, f64)>, // 90% confidence intervals
|
|
pub attention_weights: HashMap<String, Vec<f64>>, // Attention interpretability
|
|
pub feature_importance: Vec<f64>, // Variable importance scores
|
|
pub latency_us: u64, // Inference latency
|
|
}
|
|
|
|
/// Complete Temporal Fusion Transformer
|
|
pub struct TemporalFusionTransformer {
|
|
pub config: TFTConfig,
|
|
pub metadata: TFTMetadata,
|
|
pub is_trained: bool,
|
|
|
|
// Core TFT components
|
|
static_variable_selection: VariableSelectionNetwork,
|
|
historical_variable_selection: VariableSelectionNetwork,
|
|
future_variable_selection: VariableSelectionNetwork,
|
|
|
|
// Encoding layers
|
|
static_encoder: GRNStack,
|
|
historical_encoder: GRNStack,
|
|
future_encoder: GRNStack,
|
|
|
|
// Temporal processing
|
|
lstm_encoder: Linear, // Simplified LSTM representation
|
|
lstm_decoder: Linear,
|
|
|
|
// Attention mechanism
|
|
temporal_attention: TemporalSelfAttention,
|
|
|
|
// Output layers
|
|
quantile_outputs: QuantileLayer,
|
|
|
|
// Performance tracking
|
|
inference_count: AtomicU64,
|
|
total_latency_us: AtomicU64,
|
|
max_latency_us: AtomicU64,
|
|
|
|
device: Device,
|
|
|
|
// Variable map for checkpointing
|
|
varmap: Arc<VarMap>,
|
|
}
|
|
|
|
impl std::fmt::Debug for TemporalFusionTransformer {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("TemporalFusionTransformer")
|
|
.field("config", &self.config)
|
|
.field("metadata", &self.metadata)
|
|
.field("is_trained", &self.is_trained)
|
|
.field(
|
|
"inference_count",
|
|
&self.inference_count.load(Ordering::Relaxed),
|
|
)
|
|
.field(
|
|
"total_latency_us",
|
|
&self.total_latency_us.load(Ordering::Relaxed),
|
|
)
|
|
.field(
|
|
"max_latency_us",
|
|
&self.max_latency_us.load(Ordering::Relaxed),
|
|
)
|
|
.field("device", &format!("{:?}", self.device))
|
|
.field("varmap", &"Arc<VarMap>")
|
|
.finish_non_exhaustive()
|
|
}
|
|
}
|
|
|
|
impl TemporalFusionTransformer {
|
|
pub fn new(config: TFTConfig) -> Result<Self, MLError> {
|
|
Self::new_with_device(config, Device::cuda_if_available(0).unwrap_or(Device::Cpu))
|
|
}
|
|
|
|
pub fn new_with_device(config: TFTConfig, device: Device) -> Result<Self, MLError> {
|
|
// Validate configuration
|
|
let total_features =
|
|
config.num_static_features + config.num_known_features + config.num_unknown_features;
|
|
if total_features != config.input_dim {
|
|
return Err(MLError::ConfigError {
|
|
reason: format!(
|
|
"Feature count mismatch: static({}) + known({}) + unknown({}) = {} != input_dim({})",
|
|
config.num_static_features,
|
|
config.num_known_features,
|
|
config.num_unknown_features,
|
|
total_features,
|
|
config.input_dim
|
|
)
|
|
});
|
|
}
|
|
|
|
// Log configuration for debugging
|
|
debug!(
|
|
"Creating TFT with {} input features (static: {}, known: {}, unknown: {})",
|
|
config.input_dim,
|
|
config.num_static_features,
|
|
config.num_known_features,
|
|
config.num_unknown_features
|
|
);
|
|
|
|
let varmap = Arc::new(VarMap::new());
|
|
let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
|
|
|
// Create variable selection networks
|
|
let static_variable_selection = VariableSelectionNetwork::new(
|
|
config.num_static_features,
|
|
config.hidden_dim,
|
|
vs.pp("static_vsn"),
|
|
)?;
|
|
|
|
let historical_variable_selection = VariableSelectionNetwork::new(
|
|
config.num_unknown_features,
|
|
config.hidden_dim,
|
|
vs.pp("historical_vsn"),
|
|
)?;
|
|
|
|
let future_variable_selection = VariableSelectionNetwork::new(
|
|
config.num_known_features,
|
|
config.hidden_dim,
|
|
vs.pp("future_vsn"),
|
|
)?;
|
|
|
|
// Create encoding stacks
|
|
let static_encoder = GRNStack::new(
|
|
config.hidden_dim,
|
|
config.hidden_dim,
|
|
config.hidden_dim,
|
|
config.num_layers,
|
|
vs.pp("static_encoder"),
|
|
)?;
|
|
|
|
let historical_encoder = GRNStack::new(
|
|
config.hidden_dim,
|
|
config.hidden_dim,
|
|
config.hidden_dim,
|
|
config.num_layers,
|
|
vs.pp("historical_encoder"),
|
|
)?;
|
|
|
|
let future_encoder = GRNStack::new(
|
|
config.hidden_dim,
|
|
config.hidden_dim,
|
|
config.hidden_dim,
|
|
config.num_layers,
|
|
vs.pp("future_encoder"),
|
|
)?;
|
|
|
|
// Simplified LSTM layers (in practice, would use proper LSTM)
|
|
let lstm_encoder = linear(config.hidden_dim, config.hidden_dim, vs.pp("lstm_encoder"))?;
|
|
let lstm_decoder = linear(config.hidden_dim, config.hidden_dim, vs.pp("lstm_decoder"))?;
|
|
|
|
// Temporal attention
|
|
let temporal_attention = TemporalSelfAttention::new(
|
|
config.hidden_dim,
|
|
config.num_heads,
|
|
config.dropout_rate,
|
|
config.use_flash_attention,
|
|
vs.pp("temporal_attention"),
|
|
)?;
|
|
|
|
// Quantile output layer
|
|
let quantile_outputs = QuantileLayer::new(
|
|
config.hidden_dim,
|
|
config.prediction_horizon,
|
|
config.num_quantiles,
|
|
vs.pp("quantile_outputs"),
|
|
)?;
|
|
|
|
// Metadata
|
|
let metadata = TFTMetadata {
|
|
model_id: Uuid::new_v4().to_string(),
|
|
version: "1.0.0".to_string(),
|
|
input_dim: config.input_dim,
|
|
output_dim: config.prediction_horizon,
|
|
created_at: SystemTime::now(),
|
|
last_trained: None,
|
|
training_samples: 0,
|
|
performance_metrics: HashMap::new(),
|
|
};
|
|
|
|
Ok(Self {
|
|
config,
|
|
metadata,
|
|
is_trained: false,
|
|
static_variable_selection,
|
|
historical_variable_selection,
|
|
future_variable_selection,
|
|
static_encoder,
|
|
historical_encoder,
|
|
future_encoder,
|
|
lstm_encoder,
|
|
lstm_decoder,
|
|
temporal_attention,
|
|
quantile_outputs,
|
|
inference_count: AtomicU64::new(0),
|
|
total_latency_us: AtomicU64::new(0),
|
|
max_latency_us: AtomicU64::new(0),
|
|
device,
|
|
varmap,
|
|
})
|
|
}
|
|
|
|
/// Get reference to the model's VarMap for checkpointing and quantization
|
|
pub fn varmap(&self) -> &Arc<VarMap> {
|
|
&self.varmap
|
|
}
|
|
|
|
/// Get reference to the model's Device
|
|
pub fn device(&self) -> &Device {
|
|
&self.device
|
|
}
|
|
|
|
/// Validate input tensor dimensions match configuration
|
|
fn validate_input_dimensions(
|
|
&self,
|
|
static_features: &Tensor,
|
|
historical_features: &Tensor,
|
|
future_features: &Tensor,
|
|
) -> Result<(), MLError> {
|
|
// Validate static features: [batch, num_static_features]
|
|
let static_dims = static_features.dims();
|
|
if static_dims.len() != 2 {
|
|
return Err(MLError::ModelError(format!(
|
|
"Static features must be 2D [batch, features], got {} dimensions",
|
|
static_dims.len()
|
|
)));
|
|
}
|
|
if static_dims[1] != self.config.num_static_features {
|
|
return Err(MLError::ModelError(format!(
|
|
"Static features dimension mismatch: expected {}, got {}",
|
|
self.config.num_static_features, static_dims[1]
|
|
)));
|
|
}
|
|
|
|
// Validate historical features: [batch, seq_len, num_unknown_features]
|
|
let hist_dims = historical_features.dims();
|
|
if hist_dims.len() != 3 {
|
|
return Err(MLError::ModelError(format!(
|
|
"Historical features must be 3D [batch, seq, features], got {} dimensions",
|
|
hist_dims.len()
|
|
)));
|
|
}
|
|
if hist_dims[2] != self.config.num_unknown_features {
|
|
return Err(MLError::ModelError(format!(
|
|
"Historical features dimension mismatch: expected {}, got {} (Wave C+D requires 210 features)",
|
|
self.config.num_unknown_features,
|
|
hist_dims[2]
|
|
)));
|
|
}
|
|
|
|
// Validate future features: [batch, horizon, num_known_features]
|
|
let fut_dims = future_features.dims();
|
|
if fut_dims.len() != 3 {
|
|
return Err(MLError::ModelError(format!(
|
|
"Future features must be 3D [batch, horizon, features], got {} dimensions",
|
|
fut_dims.len()
|
|
)));
|
|
}
|
|
if fut_dims[2] != self.config.num_known_features {
|
|
return Err(MLError::ModelError(format!(
|
|
"Future features dimension mismatch: expected {}, got {}",
|
|
self.config.num_known_features, fut_dims[2]
|
|
)));
|
|
}
|
|
|
|
// Verify total feature count matches 225 (Wave C+D)
|
|
let total_features = self.config.num_static_features
|
|
+ self.config.num_unknown_features
|
|
+ self.config.num_known_features;
|
|
if total_features != 225 {
|
|
warn!(
|
|
"TFT configured with {} features, expected 225 for Wave C+D compatibility",
|
|
total_features
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Forward pass through the complete `TFT` architecture
|
|
#[instrument(skip(self, static_features, historical_features, future_features))]
|
|
pub fn forward(
|
|
&mut self,
|
|
static_features: &Tensor,
|
|
historical_features: &Tensor,
|
|
future_features: &Tensor,
|
|
) -> Result<Tensor, MLError> {
|
|
self.forward_with_checkpointing(
|
|
static_features,
|
|
historical_features,
|
|
future_features,
|
|
false,
|
|
)
|
|
}
|
|
|
|
/// Forward pass with optional gradient checkpointing
|
|
///
|
|
/// When gradient checkpointing is enabled:
|
|
/// - Memory usage reduced by 30-40% (doesn't store intermediate activations)
|
|
/// - Training time increases by ~20% (recomputes activations during backprop)
|
|
///
|
|
/// # Arguments
|
|
/// * `static_features` - Static input features [batch, num_static_features]
|
|
/// * `historical_features` - Historical features [batch, seq_len, num_unknown_features]
|
|
/// * `future_features` - Future features [batch, horizon, num_known_features]
|
|
/// * `use_checkpointing` - Whether to use gradient checkpointing
|
|
#[instrument(skip(self, static_features, historical_features, future_features))]
|
|
pub fn forward_with_checkpointing(
|
|
&mut self,
|
|
static_features: &Tensor,
|
|
historical_features: &Tensor,
|
|
future_features: &Tensor,
|
|
use_checkpointing: bool,
|
|
) -> Result<Tensor, MLError> {
|
|
let start_time = Instant::now();
|
|
|
|
// Validate input dimensions
|
|
self.validate_input_dimensions(static_features, historical_features, future_features)?;
|
|
|
|
// Log device placement for debugging
|
|
debug!("Forward pass device check:");
|
|
debug!(" static_features: {:?}", static_features.device());
|
|
debug!(" historical_features: {:?}", historical_features.device());
|
|
debug!(" future_features: {:?}", future_features.device());
|
|
debug!(" model device: {:?}", self.device);
|
|
|
|
// 1. Variable Selection Networks
|
|
// CRITICAL: Add .to_device() to ensure GPU execution
|
|
let static_selected = self
|
|
.static_variable_selection
|
|
.forward(static_features, None)?;
|
|
let historical_selected = self
|
|
.historical_variable_selection
|
|
.forward(historical_features, None)?;
|
|
let future_selected = self
|
|
.future_variable_selection
|
|
.forward(future_features, None)?;
|
|
|
|
debug!(" static_selected: {:?}", static_selected.device());
|
|
debug!(" historical_selected: {:?}", historical_selected.device());
|
|
debug!(" future_selected: {:?}", future_selected.device());
|
|
|
|
// 2. Feature Encoding (checkpoint expensive layers)
|
|
// CRITICAL: Add .to_device() to ensure GPU execution
|
|
let static_encoded = if use_checkpointing {
|
|
// Detach intermediate tensors to free memory during forward pass
|
|
// They will be recomputed during backward pass
|
|
self.static_encoder
|
|
.forward(&static_selected.detach(), None)?
|
|
} else {
|
|
self.static_encoder.forward(&static_selected, None)?
|
|
};
|
|
|
|
let historical_encoded = if use_checkpointing {
|
|
self.historical_encoder
|
|
.forward(&historical_selected.detach(), None)?
|
|
} else {
|
|
self.historical_encoder
|
|
.forward(&historical_selected, None)?
|
|
};
|
|
|
|
let future_encoded = if use_checkpointing {
|
|
self.future_encoder
|
|
.forward(&future_selected.detach(), None)?
|
|
} else {
|
|
self.future_encoder.forward(&future_selected, None)?
|
|
};
|
|
|
|
debug!(" static_encoded: {:?}", static_encoded.device());
|
|
debug!(" historical_encoded: {:?}", historical_encoded.device());
|
|
debug!(" future_encoded: {:?}", future_encoded.device());
|
|
|
|
// 3. Temporal Processing (checkpoint LSTM layers - most memory intensive)
|
|
// CRITICAL: Add .to_device() to ensure GPU execution
|
|
let historical_temporal = if use_checkpointing {
|
|
self.lstm_encoder.forward(&historical_encoded.detach())?
|
|
} else {
|
|
self.lstm_encoder.forward(&historical_encoded)?
|
|
};
|
|
|
|
let future_temporal = if use_checkpointing {
|
|
self.lstm_decoder.forward(&future_encoded.detach())?
|
|
} else {
|
|
self.lstm_decoder.forward(&future_encoded)?
|
|
};
|
|
|
|
debug!(" historical_temporal: {:?}", historical_temporal.device());
|
|
debug!(" future_temporal: {:?}", future_temporal.device());
|
|
|
|
// 4. Combine temporal representations
|
|
let combined_temporal =
|
|
self.combine_temporal_features(&historical_temporal, &future_temporal)?;
|
|
|
|
debug!(" combined_temporal: {:?}", combined_temporal.device());
|
|
|
|
// 5. Self-Attention (checkpoint attention - memory intensive)
|
|
// CRITICAL: Add .to_device() to ensure GPU execution
|
|
// Use specialized attention checkpointing for maximum memory savings
|
|
let attended = self.temporal_attention.forward_with_checkpointing(
|
|
&combined_temporal,
|
|
true,
|
|
use_checkpointing,
|
|
)?;
|
|
|
|
debug!(" attended: {:?}", attended.device());
|
|
|
|
// 6. Final processing with static context
|
|
let contextualized = self.apply_static_context(&attended, &static_encoded)?;
|
|
|
|
debug!(" contextualized: {:?}", contextualized.device());
|
|
|
|
// 7. Quantile Outputs (no checkpointing on final layer)
|
|
let quantile_preds = self.quantile_outputs.forward(&contextualized)?;
|
|
|
|
debug!(" quantile_preds: {:?}", quantile_preds.device());
|
|
|
|
// Update performance metrics
|
|
let latency = start_time.elapsed().as_micros() as u64;
|
|
self.update_performance_metrics(latency);
|
|
|
|
Ok(quantile_preds)
|
|
}
|
|
|
|
fn combine_temporal_features(
|
|
&self,
|
|
historical: &Tensor,
|
|
future: &Tensor,
|
|
) -> Result<Tensor, MLError> {
|
|
// Concatenate historical and future features along the time dimension
|
|
let combined = Tensor::cat(&[historical, future], 1)?;
|
|
Ok(combined)
|
|
}
|
|
|
|
fn apply_static_context(
|
|
&self,
|
|
temporal: &Tensor,
|
|
static_context: &Tensor,
|
|
) -> Result<Tensor, MLError> {
|
|
let (batch_size, seq_len, hidden_dim) = temporal.dims3()?;
|
|
|
|
// Static context comes from variable selection + GRN encoding
|
|
// It has shape [batch, 1, hidden] (variable selection adds seq_len=1 dimension)
|
|
// We need to expand it to [batch, seq_len, hidden] to match temporal features
|
|
|
|
// First, squeeze out the seq_len=1 dimension to get [batch, hidden]
|
|
let static_squeezed = static_context.squeeze(1)?;
|
|
|
|
// Then expand to match sequence length using broadcast (zero-copy)
|
|
let static_expanded = static_squeezed
|
|
.unsqueeze(1)? // [batch, 1, hidden]
|
|
.broadcast_as((batch_size, seq_len, hidden_dim))?; // [batch, seq_len, hidden] - zero-copy broadcast
|
|
|
|
// Add static context to temporal features
|
|
let contextualized = (temporal + &static_expanded)?;
|
|
|
|
Ok(contextualized)
|
|
}
|
|
|
|
/// Multi-horizon prediction interface
|
|
pub fn predict_horizons(
|
|
&mut self,
|
|
static_features: &Array1<f64>,
|
|
historical_features: &Array2<f64>,
|
|
future_features: &Array2<f64>,
|
|
) -> Result<MultiHorizonPrediction, MLError> {
|
|
if !self.is_trained {
|
|
return Err(MLError::ModelError("Model not trained".to_string()));
|
|
}
|
|
|
|
let start_time = Instant::now();
|
|
|
|
// Convert ndarray to tensors
|
|
let static_tensor = self.array_to_tensor_1d(static_features)?;
|
|
let historical_tensor = self.array_to_tensor_2d(historical_features)?;
|
|
let future_tensor = self.array_to_tensor_2d(future_features)?;
|
|
|
|
// Add batch dimension
|
|
let static_batched = static_tensor.unsqueeze(0)?;
|
|
let historical_batched = historical_tensor.unsqueeze(0)?;
|
|
let future_batched = future_tensor.unsqueeze(0)?;
|
|
|
|
// Forward pass
|
|
let quantile_preds = self.forward(&static_batched, &historical_batched, &future_batched)?;
|
|
|
|
// Extract predictions and process outputs
|
|
let pred_data = quantile_preds.squeeze(0)?.to_vec2::<f32>()?; // [horizon, quantiles]
|
|
|
|
let mut predictions = Vec::new();
|
|
let mut quantiles = Vec::new();
|
|
let mut uncertainty = Vec::new();
|
|
let mut confidence_intervals = Vec::new();
|
|
|
|
for horizon in 0..self.config.prediction_horizon {
|
|
let horizon_quantiles = &pred_data[horizon];
|
|
|
|
// Point prediction (median)
|
|
let median_idx = self.config.num_quantiles / 2;
|
|
predictions.push(horizon_quantiles[median_idx] as f64);
|
|
|
|
// All quantiles for this horizon
|
|
quantiles.push(horizon_quantiles.iter().map(|&x| x as f64).collect());
|
|
|
|
// Uncertainty (IQR)
|
|
let q75_idx = (self.config.num_quantiles * 3) / 4;
|
|
let q25_idx = self.config.num_quantiles / 4;
|
|
let iqr = horizon_quantiles[q75_idx] - horizon_quantiles[q25_idx];
|
|
uncertainty.push(iqr as f64);
|
|
|
|
// 90% confidence interval
|
|
let lower_idx = self.config.num_quantiles / 10; // ~10th percentile
|
|
let upper_idx = (self.config.num_quantiles * 9) / 10; // ~90th percentile
|
|
let ci = (
|
|
horizon_quantiles[lower_idx] as f64,
|
|
horizon_quantiles[upper_idx] as f64,
|
|
);
|
|
confidence_intervals.push(ci);
|
|
}
|
|
|
|
// Get feature importance and attention weights
|
|
let feature_importance = self.static_variable_selection.get_importance_scores()?;
|
|
let mut attention_weights = HashMap::new();
|
|
let weights = self.temporal_attention.get_attention_weights();
|
|
for (key, weight) in weights {
|
|
attention_weights.insert(key, vec![weight]);
|
|
}
|
|
|
|
let latency = start_time.elapsed().as_micros() as u64;
|
|
|
|
Ok(MultiHorizonPrediction {
|
|
predictions,
|
|
quantiles,
|
|
uncertainty,
|
|
confidence_intervals,
|
|
attention_weights,
|
|
feature_importance,
|
|
latency_us: latency,
|
|
})
|
|
}
|
|
|
|
fn array_to_tensor_1d(&self, arr: &Array1<f64>) -> Result<Tensor, MLError> {
|
|
let data: Vec<f32> = arr.iter().map(|&x| x as f32).collect();
|
|
let tensor = Tensor::from_slice(&data, arr.len(), &self.device)?;
|
|
Ok(tensor)
|
|
}
|
|
|
|
fn array_to_tensor_2d(&self, arr: &Array2<f64>) -> Result<Tensor, MLError> {
|
|
let data: Vec<f32> = arr.iter().map(|&x| x as f32).collect();
|
|
let shape = arr.shape();
|
|
let tensor = Tensor::from_slice(&data, (shape[0], shape[1]), &self.device)?;
|
|
Ok(tensor)
|
|
}
|
|
|
|
fn update_performance_metrics(&self, latency_us: u64) {
|
|
self.inference_count.fetch_add(1, Ordering::Relaxed);
|
|
self.total_latency_us
|
|
.fetch_add(latency_us, Ordering::Relaxed);
|
|
|
|
// Update max latency atomically
|
|
let mut current_max = self.max_latency_us.load(Ordering::Relaxed);
|
|
while latency_us > current_max {
|
|
match self.max_latency_us.compare_exchange_weak(
|
|
current_max,
|
|
latency_us,
|
|
Ordering::Relaxed,
|
|
Ordering::Relaxed,
|
|
) {
|
|
Ok(_) => break,
|
|
Err(new_max) => current_max = new_max,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Get performance metrics
|
|
pub fn get_metrics(&self) -> HashMap<String, f64> {
|
|
let inference_count = self.inference_count.load(Ordering::Relaxed);
|
|
let total_latency = self.total_latency_us.load(Ordering::Relaxed);
|
|
let max_latency = self.max_latency_us.load(Ordering::Relaxed);
|
|
|
|
let avg_latency = if inference_count > 0 {
|
|
total_latency as f64 / inference_count as f64
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
let throughput = if avg_latency > 0.0 {
|
|
1_000_000.0 / avg_latency // predictions per second
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
let mut metrics = HashMap::new();
|
|
metrics.insert("total_inferences".to_string(), inference_count as f64);
|
|
metrics.insert("avg_latency_us".to_string(), avg_latency);
|
|
metrics.insert("max_latency_us".to_string(), max_latency as f64);
|
|
metrics.insert("throughput_pps".to_string(), throughput);
|
|
|
|
metrics
|
|
}
|
|
|
|
/// Get reference to VarMap for weight extraction
|
|
pub fn get_varmap(&self) -> &Arc<VarMap> {
|
|
&self.varmap
|
|
}
|
|
|
|
/// Training interface with real backward pass and optimizer
|
|
pub async fn train(
|
|
&mut self,
|
|
training_data: &[(Array1<f64>, Array2<f64>, Array2<f64>, Array1<f64>)], // (static, historical, future, targets)
|
|
validation_data: &[(Array1<f64>, Array2<f64>, Array2<f64>, Array1<f64>)],
|
|
epochs: usize,
|
|
) -> Result<(), MLError> {
|
|
info!("Starting TFT training for {} epochs", epochs);
|
|
|
|
// Initialize AdamW optimizer with model parameters
|
|
let params = self.varmap.all_vars();
|
|
let num_params: usize = params.iter().map(|v| v.as_tensor().elem_count()).sum();
|
|
let lr = 1e-3;
|
|
let mut optimizer = AdamW::new(
|
|
params,
|
|
ParamsAdamW {
|
|
lr,
|
|
beta1: 0.9,
|
|
beta2: 0.999,
|
|
eps: 1e-8,
|
|
weight_decay: 1e-4,
|
|
},
|
|
)
|
|
.map_err(|e| MLError::TrainingError(format!("Failed to create optimizer: {}", e)))?;
|
|
|
|
info!(
|
|
"Initialized AdamW optimizer: lr={:.2e}, {} parameters",
|
|
lr, num_params
|
|
);
|
|
|
|
let mut best_val_loss = f64::MAX;
|
|
|
|
for epoch in 0..epochs {
|
|
let mut epoch_loss = 0.0;
|
|
|
|
for (_i, (static_feat, hist_feat, fut_feat, targets)) in
|
|
training_data.iter().enumerate()
|
|
{
|
|
// Convert to tensors
|
|
let static_tensor = self.array_to_tensor_1d(static_feat)?.unsqueeze(0)?;
|
|
let hist_tensor = self.array_to_tensor_2d(hist_feat)?.unsqueeze(0)?;
|
|
let fut_tensor = self.array_to_tensor_2d(fut_feat)?.unsqueeze(0)?;
|
|
let target_tensor = self.array_to_tensor_1d(targets)?.unsqueeze(0)?;
|
|
|
|
// Forward pass
|
|
let predictions = self.forward(&static_tensor, &hist_tensor, &fut_tensor)?;
|
|
|
|
// Compute quantile loss
|
|
let loss = self
|
|
.quantile_outputs
|
|
.quantile_loss(&predictions, &target_tensor)?;
|
|
epoch_loss += loss.to_vec0::<f32>()? as f64;
|
|
|
|
// Backward pass — compute gradients
|
|
let grads = loss.backward().map_err(|e| {
|
|
MLError::TrainingError(format!("Backward pass failed: {}", e))
|
|
})?;
|
|
|
|
// Check gradient health before stepping
|
|
let varmap_data = self.varmap.data().lock().map_err(|e| {
|
|
MLError::TrainingError(format!("Failed to lock VarMap: {}", e))
|
|
})?;
|
|
let mut grad_norm_sq = 0.0_f64;
|
|
for (_name, var) in varmap_data.iter() {
|
|
if let Some(grad) = grads.get(var.as_tensor()) {
|
|
let norm_sq = grad
|
|
.sqr()
|
|
.and_then(|t| t.sum_all())
|
|
.and_then(|t| t.to_dtype(candle_core::DType::F64))
|
|
.and_then(|t| t.to_scalar::<f64>())
|
|
.unwrap_or(0.0);
|
|
grad_norm_sq += norm_sq;
|
|
}
|
|
}
|
|
drop(varmap_data);
|
|
let grad_norm = grad_norm_sq.sqrt();
|
|
|
|
if grad_norm.is_nan() || grad_norm.is_infinite() {
|
|
warn!(
|
|
"Gradient explosion detected (norm={}), skipping update",
|
|
grad_norm
|
|
);
|
|
continue;
|
|
}
|
|
|
|
// Optimizer step — update parameters
|
|
optimizer.step(&grads).map_err(|e| {
|
|
MLError::TrainingError(format!("Optimizer step failed: {}", e))
|
|
})?;
|
|
}
|
|
|
|
let avg_epoch_loss = epoch_loss / training_data.len().max(1) as f64;
|
|
debug!("Epoch {}: Average Loss = {:.6}", epoch, avg_epoch_loss);
|
|
|
|
// Validation every 10 epochs
|
|
if epoch % 10 == 0 {
|
|
let val_loss = self.validate(validation_data).await?;
|
|
info!(
|
|
"Epoch {}: Train Loss = {:.6}, Val Loss = {:.6}",
|
|
epoch, avg_epoch_loss, val_loss
|
|
);
|
|
if val_loss < best_val_loss {
|
|
best_val_loss = val_loss;
|
|
}
|
|
}
|
|
}
|
|
|
|
self.is_trained = true;
|
|
self.metadata.last_trained = Some(SystemTime::now());
|
|
self.metadata.training_samples = training_data.len() as u64;
|
|
|
|
info!(
|
|
"TFT training completed: {} epochs, best val loss = {:.6}",
|
|
epochs, best_val_loss
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
async fn validate(
|
|
&mut self,
|
|
validation_data: &[(Array1<f64>, Array2<f64>, Array2<f64>, Array1<f64>)],
|
|
) -> Result<f64, MLError> {
|
|
let mut total_loss = 0.0;
|
|
|
|
for (static_feat, hist_feat, fut_feat, targets) in validation_data {
|
|
let static_tensor = self.array_to_tensor_1d(static_feat)?.unsqueeze(0)?;
|
|
let hist_tensor = self.array_to_tensor_2d(hist_feat)?.unsqueeze(0)?;
|
|
let fut_tensor = self.array_to_tensor_2d(fut_feat)?.unsqueeze(0)?;
|
|
let target_tensor = self.array_to_tensor_1d(targets)?.unsqueeze(0)?;
|
|
|
|
let predictions = self.forward(&static_tensor, &hist_tensor, &fut_tensor)?;
|
|
let loss = self
|
|
.quantile_outputs
|
|
.quantile_loss(&predictions, &target_tensor)?;
|
|
total_loss += loss.to_vec0::<f32>()? as f64;
|
|
}
|
|
|
|
Ok(total_loss / validation_data.len() as f64)
|
|
}
|
|
|
|
/// Compute quantile loss for training
|
|
pub fn compute_quantile_loss(
|
|
&self,
|
|
predictions: &Tensor,
|
|
targets: &Tensor,
|
|
) -> Result<Tensor, MLError> {
|
|
self.quantile_outputs.quantile_loss(predictions, targets)
|
|
}
|
|
|
|
/// HFT-optimized inference
|
|
pub fn predict_fast(
|
|
&mut self,
|
|
static_features: &[f32],
|
|
historical_features: &[f32],
|
|
future_features: &[f32],
|
|
) -> Result<Vec<f32>, MLError> {
|
|
let start = Instant::now();
|
|
|
|
// Convert to tensors (optimized path)
|
|
let static_tensor =
|
|
Tensor::from_slice(static_features, static_features.len(), &self.device)?
|
|
.unsqueeze(0)?;
|
|
|
|
let hist_len = self.config.sequence_length;
|
|
let hist_dim = self.config.num_unknown_features;
|
|
let historical_tensor =
|
|
Tensor::from_slice(historical_features, (hist_len, hist_dim), &self.device)?
|
|
.unsqueeze(0)?;
|
|
|
|
let fut_len = self.config.prediction_horizon;
|
|
let fut_dim = self.config.num_known_features;
|
|
let future_tensor =
|
|
Tensor::from_slice(future_features, (fut_len, fut_dim), &self.device)?.unsqueeze(0)?;
|
|
|
|
// Forward pass
|
|
let quantile_preds = self.forward(&static_tensor, &historical_tensor, &future_tensor)?;
|
|
|
|
// Extract median predictions
|
|
let pred_data = quantile_preds.squeeze(0)?.to_vec2::<f32>()?;
|
|
let median_idx = self.config.num_quantiles / 2;
|
|
let predictions: Vec<f32> = pred_data
|
|
.iter()
|
|
.map(|horizon_quantiles| horizon_quantiles[median_idx])
|
|
.collect();
|
|
|
|
let latency = start.elapsed().as_micros() as u64;
|
|
self.update_performance_metrics(latency);
|
|
|
|
if latency > self.config.max_inference_latency_us {
|
|
warn!(
|
|
"Inference latency {}μs exceeds target {}μs",
|
|
latency, self.config.max_inference_latency_us
|
|
);
|
|
}
|
|
|
|
Ok(predictions)
|
|
}
|
|
}
|
|
|
|
/// Implement Checkpointable trait for TFT
|
|
#[async_trait]
|
|
impl Checkpointable for TemporalFusionTransformer {
|
|
fn model_type(&self) -> ModelType {
|
|
ModelType::TFT
|
|
}
|
|
|
|
fn model_name(&self) -> &str {
|
|
&self.metadata.model_id
|
|
}
|
|
|
|
fn model_version(&self) -> &str {
|
|
&self.metadata.version
|
|
}
|
|
|
|
async fn serialize_state(&self) -> Result<Vec<u8>, MLError> {
|
|
// Save VarMap to temporary file, then read as bytes
|
|
// VarMap.save() requires a Path, not a writer
|
|
let temp_dir = std::env::temp_dir();
|
|
let temp_path = temp_dir.join(format!("tft_checkpoint_{}.safetensors", Uuid::new_v4()));
|
|
|
|
// Convert temp_path to string for VarMap::save()
|
|
let temp_path_str = temp_path
|
|
.to_str()
|
|
.ok_or_else(|| MLError::ModelError("Invalid temp path".to_string()))?;
|
|
|
|
self.varmap
|
|
.save(temp_path_str)
|
|
.map_err(|e| MLError::ModelError(format!("Failed to serialize TFT state: {}", e)))?;
|
|
|
|
// Read the file into bytes
|
|
let buffer = std::fs::read(&temp_path)
|
|
.map_err(|e| MLError::ModelError(format!("Failed to read checkpoint file: {}", e)))?;
|
|
|
|
// Clean up temp file
|
|
let _ = std::fs::remove_file(&temp_path);
|
|
|
|
debug!("Serialized TFT state: {} bytes", buffer.len());
|
|
Ok(buffer)
|
|
}
|
|
|
|
async fn deserialize_state(&mut self, data: &[u8]) -> Result<(), MLError> {
|
|
// Write bytes to temporary file, then load VarMap
|
|
let temp_dir = std::env::temp_dir();
|
|
let temp_path = temp_dir.join(format!("tft_restore_{}.safetensors", Uuid::new_v4()));
|
|
|
|
std::fs::write(&temp_path, data)
|
|
.map_err(|e| MLError::ModelError(format!("Failed to write temp checkpoint: {}", e)))?;
|
|
|
|
// Convert temp_path to string for VarMap::load()
|
|
let temp_path_str = temp_path
|
|
.to_str()
|
|
.ok_or_else(|| MLError::ModelError("Invalid temp path".to_string()))?;
|
|
|
|
// Try to get mutable access to the VarMap through Arc
|
|
let varmap_mut = Arc::get_mut(&mut self.varmap).ok_or_else(|| {
|
|
MLError::ModelError(
|
|
"Cannot load checkpoint: VarMap has multiple references. \
|
|
This indicates the model is being shared across threads. \
|
|
Clone the model before loading checkpoint."
|
|
.to_string(),
|
|
)
|
|
})?;
|
|
|
|
// Load the checkpoint into the VarMap
|
|
varmap_mut
|
|
.load(temp_path_str)
|
|
.map_err(|e| MLError::ModelError(format!("Failed to load TFT state: {}", e)))?;
|
|
|
|
// Clean up temp file
|
|
let _ = std::fs::remove_file(&temp_path);
|
|
|
|
debug!("Deserialized TFT state from {} bytes", data.len());
|
|
Ok(())
|
|
}
|
|
|
|
fn get_training_state(&self) -> (Option<u64>, Option<u64>, Option<f64>, Option<f64>) {
|
|
// TFT doesn't track epochs/steps in the current implementation
|
|
// Return metadata-based info if available
|
|
(
|
|
None, // epoch
|
|
None, // step
|
|
None, // loss
|
|
None, // accuracy
|
|
)
|
|
}
|
|
|
|
fn get_hyperparameters(&self) -> HashMap<String, Value> {
|
|
let mut params = HashMap::new();
|
|
// Core architecture params (Wave C+D: 225 features)
|
|
params.insert("input_dim".to_string(), Value::from(self.config.input_dim));
|
|
params.insert(
|
|
"hidden_dim".to_string(),
|
|
Value::from(self.config.hidden_dim),
|
|
);
|
|
params.insert("num_heads".to_string(), Value::from(self.config.num_heads));
|
|
params.insert(
|
|
"num_layers".to_string(),
|
|
Value::from(self.config.num_layers),
|
|
);
|
|
params.insert(
|
|
"prediction_horizon".to_string(),
|
|
Value::from(self.config.prediction_horizon),
|
|
);
|
|
params.insert(
|
|
"sequence_length".to_string(),
|
|
Value::from(self.config.sequence_length),
|
|
);
|
|
params.insert(
|
|
"num_quantiles".to_string(),
|
|
Value::from(self.config.num_quantiles),
|
|
);
|
|
|
|
// Feature split (critical for Wave C+D compatibility)
|
|
params.insert(
|
|
"num_static_features".to_string(),
|
|
Value::from(self.config.num_static_features),
|
|
);
|
|
params.insert(
|
|
"num_known_features".to_string(),
|
|
Value::from(self.config.num_known_features),
|
|
);
|
|
params.insert(
|
|
"num_unknown_features".to_string(),
|
|
Value::from(self.config.num_unknown_features),
|
|
);
|
|
|
|
// Training params
|
|
params.insert(
|
|
"learning_rate".to_string(),
|
|
Value::from(self.config.learning_rate),
|
|
);
|
|
params.insert(
|
|
"batch_size".to_string(),
|
|
Value::from(self.config.batch_size),
|
|
);
|
|
params.insert(
|
|
"dropout_rate".to_string(),
|
|
Value::from(self.config.dropout_rate),
|
|
);
|
|
params.insert(
|
|
"l2_regularization".to_string(),
|
|
Value::from(self.config.l2_regularization),
|
|
);
|
|
|
|
// HFT optimization flags
|
|
params.insert(
|
|
"use_flash_attention".to_string(),
|
|
Value::from(self.config.use_flash_attention),
|
|
);
|
|
params.insert(
|
|
"mixed_precision".to_string(),
|
|
Value::from(self.config.mixed_precision),
|
|
);
|
|
params.insert(
|
|
"memory_efficient".to_string(),
|
|
Value::from(self.config.memory_efficient),
|
|
);
|
|
|
|
params
|
|
}
|
|
|
|
fn get_metrics(&self) -> HashMap<String, f64> {
|
|
// Call the existing get_metrics method from TemporalFusionTransformer
|
|
let inference_count = self.inference_count.load(Ordering::Relaxed);
|
|
let total_latency = self.total_latency_us.load(Ordering::Relaxed);
|
|
let max_latency = self.max_latency_us.load(Ordering::Relaxed);
|
|
|
|
let avg_latency = if inference_count > 0 {
|
|
total_latency as f64 / inference_count as f64
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
let throughput = if avg_latency > 0.0 {
|
|
1_000_000.0 / avg_latency
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
let mut metrics = HashMap::new();
|
|
metrics.insert("total_inferences".to_string(), inference_count as f64);
|
|
metrics.insert("avg_latency_us".to_string(), avg_latency);
|
|
metrics.insert("max_latency_us".to_string(), max_latency as f64);
|
|
metrics.insert("throughput_pps".to_string(), throughput);
|
|
metrics
|
|
}
|
|
|
|
fn get_architecture_info(&self) -> HashMap<String, Value> {
|
|
let mut info = HashMap::new();
|
|
info.insert("network_type".to_string(), Value::from("TFT"));
|
|
info.insert(
|
|
"input_dim".to_string(),
|
|
Value::from(self.metadata.input_dim),
|
|
);
|
|
info.insert(
|
|
"output_dim".to_string(),
|
|
Value::from(self.metadata.output_dim),
|
|
);
|
|
info.insert(
|
|
"hidden_dim".to_string(),
|
|
Value::from(self.config.hidden_dim),
|
|
);
|
|
info.insert("num_heads".to_string(), Value::from(self.config.num_heads));
|
|
info.insert(
|
|
"num_layers".to_string(),
|
|
Value::from(self.config.num_layers),
|
|
);
|
|
info.insert(
|
|
"num_static_features".to_string(),
|
|
Value::from(self.config.num_static_features),
|
|
);
|
|
info.insert(
|
|
"num_known_features".to_string(),
|
|
Value::from(self.config.num_known_features),
|
|
);
|
|
info.insert(
|
|
"num_unknown_features".to_string(),
|
|
Value::from(self.config.num_unknown_features),
|
|
);
|
|
info
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use anyhow::Result;
|
|
|
|
#[tokio::test]
|
|
async fn test_tft_creation() -> Result<()> {
|
|
let config = TFTConfig {
|
|
input_dim: 10,
|
|
hidden_dim: 32,
|
|
num_heads: 4,
|
|
num_quantiles: 5,
|
|
prediction_horizon: 5,
|
|
sequence_length: 20,
|
|
num_static_features: 2,
|
|
num_known_features: 3,
|
|
num_unknown_features: 5,
|
|
..Default::default()
|
|
};
|
|
|
|
let tft = TemporalFusionTransformer::new(config)
|
|
.map_err(|_| anyhow::anyhow!("Failed to create TFT"))?;
|
|
assert_eq!(tft.metadata.input_dim, 10);
|
|
assert_eq!(tft.metadata.output_dim, 5);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_tft_225_features_default() -> Result<()> {
|
|
// Test default configuration uses 225 features (Wave C+D)
|
|
let config = TFTConfig::default();
|
|
assert_eq!(
|
|
config.input_dim, 225,
|
|
"Default TFT config should use 225 features"
|
|
);
|
|
assert_eq!(config.num_static_features, 5);
|
|
assert_eq!(config.num_known_features, 10);
|
|
assert_eq!(config.num_unknown_features, 210);
|
|
|
|
let tft = TemporalFusionTransformer::new(config)
|
|
.map_err(|_| anyhow::anyhow!("Failed to create TFT with 225 features"))?;
|
|
assert_eq!(tft.metadata.input_dim, 225);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_tft_225_features_validation() -> Result<()> {
|
|
// Test that 225-feature TFT validates input dimensions correctly
|
|
let config = TFTConfig::default(); // 225 features
|
|
let device = Device::Cpu;
|
|
let tft = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())
|
|
.map_err(|_| anyhow::anyhow!("Failed to create TFT"))?;
|
|
|
|
// Create valid input tensors
|
|
let batch_size = 2;
|
|
let seq_len = 50;
|
|
let horizon = 10;
|
|
|
|
let static_features = Tensor::zeros(
|
|
(batch_size, config.num_static_features),
|
|
DType::F32,
|
|
&device,
|
|
)?;
|
|
let historical_features = Tensor::zeros(
|
|
(batch_size, seq_len, config.num_unknown_features),
|
|
DType::F32,
|
|
&device,
|
|
)?;
|
|
let future_features = Tensor::zeros(
|
|
(batch_size, horizon, config.num_known_features),
|
|
DType::F32,
|
|
&device,
|
|
)?;
|
|
|
|
// Should validate successfully
|
|
let result =
|
|
tft.validate_input_dimensions(&static_features, &historical_features, &future_features);
|
|
assert!(
|
|
result.is_ok(),
|
|
"Valid 225-feature input should pass validation"
|
|
);
|
|
|
|
// Test invalid historical features dimension
|
|
let invalid_hist = Tensor::zeros((batch_size, seq_len, 50), DType::F32, &device)?; // Wrong dim: 50 instead of 210
|
|
let result =
|
|
tft.validate_input_dimensions(&static_features, &invalid_hist, &future_features);
|
|
assert!(
|
|
result.is_err(),
|
|
"Invalid historical features should fail validation"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_tft_config_mismatch_detection() -> Result<()> {
|
|
// Test that mismatched feature counts are detected during construction
|
|
let invalid_config = TFTConfig {
|
|
input_dim: 225,
|
|
num_static_features: 5,
|
|
num_known_features: 10,
|
|
num_unknown_features: 100, // Wrong: should be 210 for 225 total
|
|
..Default::default()
|
|
};
|
|
|
|
let result = TemporalFusionTransformer::new(invalid_config);
|
|
assert!(
|
|
result.is_err(),
|
|
"Mismatched feature counts should be rejected"
|
|
);
|
|
|
|
let err_msg = format!("{:?}", result.unwrap_err());
|
|
assert!(
|
|
err_msg.contains("Feature count mismatch"),
|
|
"Error should mention feature count mismatch"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_tft_checkpoint_preserves_config() -> Result<()> {
|
|
// Test that checkpoint save/load preserves 225-feature configuration
|
|
let config = TFTConfig::default(); // 225 features
|
|
let tft = TemporalFusionTransformer::new(config.clone())
|
|
.map_err(|_| anyhow::anyhow!("Failed to create TFT"))?;
|
|
|
|
let hyperparams = tft.get_hyperparameters();
|
|
|
|
// Verify all critical config params are saved
|
|
assert_eq!(
|
|
hyperparams.get("input_dim").and_then(|v| v.as_u64()),
|
|
Some(225)
|
|
);
|
|
assert_eq!(
|
|
hyperparams
|
|
.get("num_static_features")
|
|
.and_then(|v| v.as_u64()),
|
|
Some(5)
|
|
);
|
|
assert_eq!(
|
|
hyperparams
|
|
.get("num_known_features")
|
|
.and_then(|v| v.as_u64()),
|
|
Some(10)
|
|
);
|
|
assert_eq!(
|
|
hyperparams
|
|
.get("num_unknown_features")
|
|
.and_then(|v| v.as_u64()),
|
|
Some(210)
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_tft_wave_c_config() -> Result<()> {
|
|
// Test Wave C configuration (201 features)
|
|
let wave_c_config = TFTConfig {
|
|
input_dim: 201,
|
|
num_static_features: 5,
|
|
num_known_features: 10,
|
|
num_unknown_features: 186, // 201 - 5 - 10 = 186
|
|
..Default::default()
|
|
};
|
|
|
|
let tft = TemporalFusionTransformer::new(wave_c_config)
|
|
.map_err(|_| anyhow::anyhow!("Failed to create TFT with 201 features"))?;
|
|
assert_eq!(tft.metadata.input_dim, 201);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_tft_state_creation() -> Result<()> {
|
|
let config = TFTConfig {
|
|
hidden_dim: 32,
|
|
sequence_length: 20,
|
|
num_heads: 4,
|
|
..Default::default()
|
|
};
|
|
|
|
let state =
|
|
TFTState::zeros(&config).map_err(|_| anyhow::anyhow!("Failed to create state"))?;
|
|
assert!(state.last_update == 0);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_tft_config_default() -> Result<()> {
|
|
let config = TFTConfig::default();
|
|
assert!(config.input_dim > 0);
|
|
assert!(config.hidden_dim > 0);
|
|
assert!(config.num_heads > 0);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_tft_performance_metrics() -> Result<()> {
|
|
let config = TFTConfig {
|
|
input_dim: 30,
|
|
hidden_dim: 32,
|
|
num_static_features: 5,
|
|
num_known_features: 10,
|
|
num_unknown_features: 15, // 30 - 5 - 10 = 15
|
|
..Default::default()
|
|
};
|
|
|
|
let tft = TemporalFusionTransformer::new(config)
|
|
.map_err(|_| anyhow::anyhow!("Failed to create TFT"))?;
|
|
let metrics = tft.get_metrics();
|
|
|
|
assert!(metrics.contains_key("total_inferences"));
|
|
assert!(metrics.contains_key("avg_latency_us"));
|
|
assert!(metrics.contains_key("max_latency_us"));
|
|
assert!(metrics.contains_key("throughput_pps"));
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_tft_training_state() -> Result<()> {
|
|
let config = TFTConfig::default();
|
|
let mut tft = TemporalFusionTransformer::new(config)
|
|
.map_err(|_| anyhow::anyhow!("Failed to create TFT"))?;
|
|
|
|
assert!(!tft.is_trained);
|
|
tft.is_trained = true;
|
|
assert!(tft.is_trained);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_tft_metadata() -> Result<()> {
|
|
let config = TFTConfig {
|
|
input_dim: 30,
|
|
prediction_horizon: 12,
|
|
num_static_features: 5,
|
|
num_known_features: 10,
|
|
num_unknown_features: 15, // 30 - 5 - 10 = 15
|
|
..Default::default()
|
|
};
|
|
|
|
let tft = TemporalFusionTransformer::new(config)
|
|
.map_err(|_| anyhow::anyhow!("Failed to create TFT"))?;
|
|
assert_eq!(tft.metadata.input_dim, 30);
|
|
assert_eq!(tft.metadata.output_dim, 12);
|
|
Ok(())
|
|
}
|
|
}
|