5d6e79263c
tft_real_dbn_data: StreamTensor::from_vec, quantile loss returns f32 ppo_recurrent_integration: PPO::new() API, get_policy_state &[f32] test_dbn_sequence_256: to_host + manual indexing instead of .i() ops ppo_checkpoint_roundtrip: save/load_checkpoint(&PathBuf) API mamba2_accuracy_fix: pure f64 arithmetic, no GPU tensors needed ppo_lstm_training_loop: PPO::new() API ppo_step_counter_fix: new checkpoint API ppo_recurrent_performance: forward_host, LSTM batch_size arg Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
468 lines
14 KiB
Rust
468 lines
14 KiB
Rust
#![allow(
|
|
clippy::assertions_on_constants,
|
|
clippy::assertions_on_result_states,
|
|
clippy::clone_on_copy,
|
|
clippy::decimal_literal_representation,
|
|
clippy::doc_markdown,
|
|
clippy::empty_line_after_doc_comments,
|
|
clippy::field_reassign_with_default,
|
|
clippy::get_unwrap,
|
|
clippy::identity_op,
|
|
clippy::inconsistent_digit_grouping,
|
|
clippy::indexing_slicing,
|
|
clippy::integer_division,
|
|
clippy::len_zero,
|
|
clippy::let_underscore_must_use,
|
|
clippy::manual_div_ceil,
|
|
clippy::manual_let_else,
|
|
clippy::manual_range_contains,
|
|
clippy::modulo_arithmetic,
|
|
clippy::needless_range_loop,
|
|
clippy::non_ascii_literal,
|
|
clippy::redundant_clone,
|
|
clippy::shadow_reuse,
|
|
clippy::shadow_same,
|
|
clippy::shadow_unrelated,
|
|
clippy::single_match_else,
|
|
clippy::str_to_string,
|
|
clippy::string_slice,
|
|
clippy::tests_outside_test_module,
|
|
clippy::too_many_lines,
|
|
clippy::unnecessary_wraps,
|
|
clippy::unseparated_literal_suffix,
|
|
clippy::use_debug,
|
|
clippy::useless_vec,
|
|
clippy::wildcard_enum_match_arm,
|
|
clippy::else_if_without_else,
|
|
clippy::expect_used,
|
|
clippy::missing_const_for_fn,
|
|
clippy::similar_names,
|
|
clippy::type_complexity,
|
|
clippy::collapsible_else_if,
|
|
clippy::doc_lazy_continuation,
|
|
clippy::items_after_test_module,
|
|
clippy::map_clone,
|
|
clippy::multiple_unsafe_ops_per_block,
|
|
clippy::unwrap_or_default,
|
|
clippy::assign_op_pattern,
|
|
clippy::needless_borrow,
|
|
clippy::println_empty_string,
|
|
clippy::unnecessary_cast,
|
|
clippy::used_underscore_binding,
|
|
clippy::create_dir,
|
|
clippy::implicit_saturating_sub,
|
|
clippy::exit,
|
|
clippy::expect_fun_call,
|
|
clippy::too_many_arguments,
|
|
clippy::unnecessary_map_or,
|
|
clippy::unwrap_used,
|
|
dead_code,
|
|
unused_imports,
|
|
unused_variables,
|
|
clippy::cloned_ref_to_slice_refs,
|
|
clippy::neg_multiply,
|
|
clippy::while_let_loop,
|
|
clippy::bool_assert_comparison,
|
|
clippy::excessive_precision,
|
|
clippy::trivially_copy_pass_by_ref,
|
|
clippy::op_ref,
|
|
clippy::redundant_closure,
|
|
clippy::unnecessary_lazy_evaluations,
|
|
clippy::if_then_some_else_none,
|
|
clippy::unnecessary_to_owned,
|
|
clippy::single_component_path_imports,
|
|
)]
|
|
//! Test DbnSequenceLoader produces correct 256-dimensional features
|
|
//!
|
|
//! Validates that the fixed DbnSequenceLoader correctly extracts and pads
|
|
//! features to exactly 256 dimensions for MAMBA-2 training.
|
|
|
|
use anyhow::Result;
|
|
// candle eliminated — test uses native APIs
|
|
use cudarc::driver::CudaContext;
|
|
use ml::data_loaders::DbnSequenceLoader;
|
|
use std::env;
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
use tracing::info;
|
|
use tracing::warn;
|
|
|
|
/// Get test data directory path
|
|
fn get_test_data_dir() -> PathBuf {
|
|
// Try CARGO_MANIFEST_DIR first (works in tests)
|
|
if let Ok(manifest_dir) = env::var("CARGO_MANIFEST_DIR") {
|
|
PathBuf::from(manifest_dir)
|
|
.parent()
|
|
.unwrap()
|
|
.join("test_data/real/databento/ml_training_small")
|
|
} else {
|
|
// Fallback to relative path from project root
|
|
PathBuf::from("test_data/real/databento/ml_training_small")
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_feature_dimension_256() -> Result<()> {
|
|
info!("Testing DbnSequenceLoader 256-dimensional features");
|
|
|
|
let test_dir = get_test_data_dir();
|
|
|
|
if !test_dir.exists() {
|
|
warn!(test_dir = ?test_dir, "Test data not found, skipping test");
|
|
return Ok(());
|
|
}
|
|
|
|
// Create loader with d_model=256
|
|
let mut loader = DbnSequenceLoader::with_limits(60, 256, Some(10), 10).await?;
|
|
info!("Created DbnSequenceLoader (seq_len=60, d_model=256, max=10, stride=10)");
|
|
|
|
// Load sequences
|
|
info!(test_dir = ?test_dir, "Loading sequences");
|
|
let (train_data, val_data) = loader.load_sequences(&test_dir, 0.9).await?;
|
|
|
|
let total_sequences = train_data.len() + val_data.len();
|
|
info!(
|
|
total_sequences,
|
|
train = train_data.len(),
|
|
val = val_data.len(),
|
|
"Loaded sequences"
|
|
);
|
|
|
|
// Verify at least some data was loaded
|
|
assert!(total_sequences > 0, "Should load at least some sequences");
|
|
|
|
// Test 1: Verify input tensor dimensions
|
|
info!("Test 1: Verifying input tensor dimensions");
|
|
for (idx, (input, _target)) in train_data.iter().take(5).enumerate() {
|
|
let input_dims = input.dims();
|
|
|
|
info!(idx, input_dims = ?input_dims, "Sequence input shape");
|
|
|
|
// Input should be [batch=1, seq_len=60, d_model=256]
|
|
assert_eq!(
|
|
input_dims.len(),
|
|
3,
|
|
"Input should be 3D (batch, seq_len, features), got {:?}",
|
|
input_dims
|
|
);
|
|
assert_eq!(
|
|
input_dims[0], 1,
|
|
"Batch dimension should be 1, got {}",
|
|
input_dims[0]
|
|
);
|
|
assert_eq!(
|
|
input_dims[1], 60,
|
|
"Sequence length should be 60, got {}",
|
|
input_dims[1]
|
|
);
|
|
assert_eq!(
|
|
input_dims[2], 256,
|
|
"Feature dimension should be 256, got {}",
|
|
input_dims[2]
|
|
);
|
|
}
|
|
info!("All input tensors have correct shape [1, 60, 256]");
|
|
|
|
// Test 2: Verify target tensor dimensions
|
|
info!("Test 2: Verifying target tensor dimensions");
|
|
for (idx, (_input, target)) in train_data.iter().take(5).enumerate() {
|
|
let target_dims = target.dims();
|
|
|
|
info!(idx, target_dims = ?target_dims, "Sequence target shape");
|
|
|
|
// Target should be [batch=1, timesteps=1, d_model=256]
|
|
assert_eq!(
|
|
target_dims.len(),
|
|
3,
|
|
"Target should be 3D, got {:?}",
|
|
target_dims
|
|
);
|
|
assert_eq!(
|
|
target_dims[0], 1,
|
|
"Target batch should be 1, got {}",
|
|
target_dims[0]
|
|
);
|
|
assert_eq!(
|
|
target_dims[1], 1,
|
|
"Target timesteps should be 1, got {}",
|
|
target_dims[1]
|
|
);
|
|
assert_eq!(
|
|
target_dims[2], 256,
|
|
"Target feature dim should be 256, got {}",
|
|
target_dims[2]
|
|
);
|
|
}
|
|
info!("All target tensors have correct shape [1, 1, 256]");
|
|
|
|
// Test 3: Verify feature values are normalized (not all zeros/NaN)
|
|
info!("Test 3: Verifying feature normalization");
|
|
let (first_input, _) = &train_data[0];
|
|
let stream = CudaContext::new(0)
|
|
.and_then(|ctx| ctx.new_stream())
|
|
.expect("CUDA required for readback");
|
|
let values_f32 = first_input.to_host(&stream)?;
|
|
let values: Vec<f64> = values_f32.iter().map(|&v| v as f64).collect();
|
|
|
|
// Check for NaN values
|
|
let nan_count = values.iter().filter(|v| v.is_nan()).count();
|
|
assert_eq!(nan_count, 0, "Found {} NaN values in features", nan_count);
|
|
info!("No NaN values detected");
|
|
|
|
// Check for all-zero sequences (should have some variation)
|
|
let non_zero_count = values.iter().filter(|v| v.abs() > 1e-6).count();
|
|
let non_zero_ratio = non_zero_count as f64 / values.len() as f64;
|
|
info!(
|
|
non_zero_count,
|
|
total = values.len(),
|
|
non_zero_pct = non_zero_ratio * 100.0,
|
|
"Non-zero values"
|
|
);
|
|
|
|
assert!(
|
|
non_zero_ratio > 0.01,
|
|
"Features appear to be all zeros (only {:.1}% non-zero)",
|
|
non_zero_ratio * 100.0
|
|
);
|
|
|
|
// Check value range (normalized features should be roughly in [-5, 5] range)
|
|
let min_val = values.iter().cloned().fold(f64::INFINITY, f64::min);
|
|
let max_val = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
|
|
let mean = values.iter().sum::<f64>() / values.len() as f64;
|
|
|
|
info!(min_val, max_val, mean, "Value range");
|
|
|
|
// Normalized features should not have extreme outliers
|
|
assert!(
|
|
min_val > -100.0 && max_val < 100.0,
|
|
"Feature values seem unnormalized: range [{:.2}, {:.2}]",
|
|
min_val,
|
|
max_val
|
|
);
|
|
|
|
info!("Features are properly normalized");
|
|
|
|
// Test 4: Verify validation data has same properties
|
|
info!("Test 4: Verifying validation data");
|
|
if !val_data.is_empty() {
|
|
let (val_input, val_target) = &val_data[0];
|
|
|
|
assert_eq!(
|
|
val_input.dims(),
|
|
&[1, 60, 256],
|
|
"Validation input should be [1, 60, 256]"
|
|
);
|
|
assert_eq!(
|
|
val_target.dims(),
|
|
&[1, 1, 256],
|
|
"Validation target should be [1, 1, 256]"
|
|
);
|
|
|
|
info!("Validation data shapes correct");
|
|
info!(val_count = val_data.len(), "Validation sequences verified");
|
|
} else {
|
|
warn!("No validation data (split ratio may be too high)");
|
|
}
|
|
|
|
info!(
|
|
total_sequences,
|
|
train = train_data.len(),
|
|
val = val_data.len(),
|
|
"ALL TESTS PASSED: feature_dim=256, input=[1,60,256], target=[1,1,256], normalization=valid"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_extract_features_dimension() -> Result<()> {
|
|
info!("Testing extract_features() returns 256 dimensions");
|
|
|
|
let test_dir = get_test_data_dir();
|
|
|
|
if !test_dir.exists() {
|
|
warn!("Test data not found, skipping test");
|
|
return Ok(());
|
|
}
|
|
|
|
// Create loader
|
|
let mut loader = DbnSequenceLoader::new(60, 256).await?;
|
|
info!("Created DbnSequenceLoader");
|
|
|
|
// Load sequences
|
|
let (train_data, _) = loader.load_sequences(&test_dir, 0.9).await?;
|
|
|
|
assert!(!train_data.is_empty(), "Should have training data");
|
|
|
|
// Get first sequence and verify it was created with 256-dim features
|
|
let (input, _) = &train_data[0];
|
|
|
|
// Input is [1, 60, 256] where 256 is the feature dimension
|
|
let feature_dim = input.dims()[2];
|
|
|
|
info!(feature_dim, "Feature dimension from tensor");
|
|
assert_eq!(
|
|
feature_dim, 256,
|
|
"Feature dimension should be 256, got {}",
|
|
feature_dim
|
|
);
|
|
|
|
info!("extract_features() correctly produces 256-dimensional features");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_different_d_model_values() -> Result<()> {
|
|
info!("Testing different d_model values (128, 256, 512)");
|
|
|
|
let test_dir = get_test_data_dir();
|
|
|
|
if !test_dir.exists() {
|
|
warn!("Test data not found, skipping test");
|
|
return Ok(());
|
|
}
|
|
|
|
// Test different d_model values
|
|
let d_models = vec![128, 256, 512];
|
|
|
|
for d_model in d_models {
|
|
info!(d_model, "Testing d_model");
|
|
|
|
let mut loader = DbnSequenceLoader::with_limits(60, d_model, Some(5), 10).await?;
|
|
let (train_data, _) = loader.load_sequences(&test_dir, 0.9).await?;
|
|
|
|
if !train_data.is_empty() {
|
|
let (input, target) = &train_data[0];
|
|
|
|
// Verify input shape
|
|
assert_eq!(
|
|
input.dims()[2],
|
|
d_model,
|
|
"Input feature dim should be {}",
|
|
d_model
|
|
);
|
|
|
|
// Verify target shape
|
|
assert_eq!(
|
|
target.dims()[2],
|
|
d_model,
|
|
"Target feature dim should be {}",
|
|
d_model
|
|
);
|
|
|
|
info!(
|
|
d_model,
|
|
input_dims = ?input.dims(),
|
|
target_dims = ?target.dims(),
|
|
"d_model dimensions verified"
|
|
);
|
|
}
|
|
}
|
|
|
|
info!("All d_model values produce correct dimensions");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_sequence_temporal_ordering() -> Result<()> {
|
|
info!("Testing temporal ordering of sequences");
|
|
|
|
let test_dir = get_test_data_dir();
|
|
|
|
if !test_dir.exists() {
|
|
warn!("Test data not found, skipping test");
|
|
return Ok(());
|
|
}
|
|
|
|
// Create loader with stride=1 to get consecutive sequences
|
|
let mut loader = DbnSequenceLoader::with_limits(10, 256, Some(3), 1).await?;
|
|
let (train_data, _) = loader.load_sequences(&test_dir, 0.9).await?;
|
|
|
|
if train_data.len() >= 2 {
|
|
info!("Comparing consecutive sequences");
|
|
|
|
let (seq1_input, _) = &train_data[0];
|
|
let (seq2_input, _) = &train_data[1];
|
|
|
|
// With stride=1, the second sequence should be a shifted version of the first
|
|
// seq1: [t0, t1, t2, ..., t9]
|
|
// seq2: [t1, t2, t3, ..., t10]
|
|
|
|
let stream = CudaContext::new(0)
|
|
.and_then(|ctx| ctx.new_stream())
|
|
.expect("CUDA required for readback");
|
|
|
|
// Download both to host for comparison
|
|
let seq1_host = seq1_input.to_host(&stream)?;
|
|
let seq2_host = seq2_input.to_host(&stream)?;
|
|
|
|
// seq1_input shape: [1, 10, 256], seq2_input shape: [1, 10, 256]
|
|
let d_model = seq1_input.shape()[2]; // 256
|
|
|
|
// Extract last 9 timesteps from seq1 and first 9 from seq2
|
|
// seq1[0, 1..10, :] vs seq2[0, 0..9, :]
|
|
let mut max_diff: f64 = 0.0;
|
|
for t in 0..9 {
|
|
for f in 0..d_model {
|
|
let seq1_val = seq1_host[(t + 1) * d_model + f] as f64;
|
|
let seq2_val = seq2_host[t * d_model + f] as f64;
|
|
let diff = (seq1_val - seq2_val).abs();
|
|
if diff > max_diff {
|
|
max_diff = diff;
|
|
}
|
|
}
|
|
}
|
|
|
|
info!(max_diff, "Max difference between overlapping windows");
|
|
|
|
assert!(
|
|
max_diff < 1e-6,
|
|
"Consecutive sequences should overlap with stride=1, max_diff={}",
|
|
max_diff
|
|
);
|
|
}
|
|
|
|
info!("Temporal ordering verified");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_batch_processing() -> Result<()> {
|
|
info!("Testing batch processing with multiple sequences");
|
|
|
|
let test_dir = get_test_data_dir();
|
|
|
|
if !test_dir.exists() {
|
|
warn!("Test data not found, skipping test");
|
|
return Ok(());
|
|
}
|
|
|
|
// Load multiple sequences
|
|
let mut loader = DbnSequenceLoader::with_limits(60, 256, Some(100), 10).await?;
|
|
let (train_data, val_data) = loader.load_sequences(&test_dir, 0.8).await?;
|
|
|
|
let total = train_data.len() + val_data.len();
|
|
info!(total, "Loaded sequences");
|
|
|
|
// Verify all sequences have consistent dimensions
|
|
let mut valid_count = 0;
|
|
for (input, target) in train_data.iter().chain(val_data.iter()) {
|
|
if input.dims() == &[1, 60, 256] && target.dims() == &[1, 1, 256] {
|
|
valid_count += 1;
|
|
}
|
|
}
|
|
|
|
info!(valid_count, total, "Sequences with correct dimensions");
|
|
assert_eq!(
|
|
valid_count, total,
|
|
"All sequences should have correct dimensions"
|
|
);
|
|
|
|
info!("Batch processing verified");
|
|
|
|
Ok(())
|
|
}
|