563 lines
17 KiB
Rust
563 lines
17 KiB
Rust
use anyhow::Result;
|
|
use rtx_compress::quantization::vector_quantization::{
|
|
CodebookInitialization, DistanceMetric, VQConfig, VectorQuantizer,
|
|
};
|
|
use rtx_tensor::{DType, Device, Tensor};
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal device randn issue"]
|
|
fn test_vq_codebook_initialization() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Create test data with some structure
|
|
let data = Tensor::randn(&[1000, 128], &device)?;
|
|
|
|
let config = VQConfig {
|
|
codebook_size: 256,
|
|
vector_dim: 128,
|
|
max_iterations: 100,
|
|
tolerance: 1e-6,
|
|
initialization: CodebookInitialization::KMeansPlusPlus,
|
|
distance_metric: DistanceMetric::Euclidean,
|
|
};
|
|
|
|
let mut vq = VectorQuantizer::new(config);
|
|
|
|
// Test different initialization methods
|
|
let init_methods = vec![
|
|
CodebookInitialization::Random,
|
|
CodebookInitialization::KMeansPlusPlus,
|
|
CodebookInitialization::FromData,
|
|
];
|
|
|
|
for init_method in init_methods {
|
|
let mut vq_test = VectorQuantizer::new(VQConfig {
|
|
initialization: init_method.clone(),
|
|
..vq.config().clone()
|
|
});
|
|
|
|
vq_test.fit(&data)?;
|
|
|
|
// Verify codebook properties
|
|
let codebook = vq_test.codebook();
|
|
assert_eq!(codebook.shape(), &[256, 128]);
|
|
|
|
// Test encoding
|
|
let codes = vq_test.encode(&data)?;
|
|
assert_eq!(codes.shape(), &[1000]);
|
|
|
|
// All codes should be valid indices
|
|
let max_code = codes.max()?.to_scalar::<f32>()? as usize;
|
|
// For min, negate and use max
|
|
let neg_codes = codes.mul_scalar(-1.0)?;
|
|
let min_code = (neg_codes.max()?.to_scalar::<f32>()? * -1.0) as usize;
|
|
|
|
assert!(max_code < 256, "All codes should be < codebook_size");
|
|
assert!(min_code < 256, "Min code should be valid");
|
|
|
|
println!(
|
|
"Init method {:?}: max_code={}, min_code={}",
|
|
init_method, max_code, min_code
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal device randn issue"]
|
|
fn test_vq_distance_metrics() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
let data = Tensor::randn(&[500, 64], &device)?;
|
|
|
|
let distance_metrics = vec![
|
|
DistanceMetric::Euclidean,
|
|
DistanceMetric::Cosine,
|
|
DistanceMetric::Manhattan,
|
|
];
|
|
|
|
let mut results = vec![];
|
|
|
|
for distance_metric in distance_metrics {
|
|
let config = VQConfig {
|
|
codebook_size: 128,
|
|
vector_dim: 64,
|
|
max_iterations: 50,
|
|
tolerance: 1e-5,
|
|
initialization: CodebookInitialization::KMeansPlusPlus,
|
|
distance_metric: distance_metric.clone(),
|
|
};
|
|
|
|
let mut vq = VectorQuantizer::new(config);
|
|
vq.fit(&data)?;
|
|
|
|
let codes = vq.encode(&data)?;
|
|
let reconstructed = vq.decode(&codes)?;
|
|
|
|
// Compute reconstruction error using the same distance metric
|
|
let reconstruction_error = match distance_metric {
|
|
DistanceMetric::Euclidean => {
|
|
let diff = (data.clone() - reconstructed.clone())?;
|
|
let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
|
|
let squared = diff.pow(&two)?;
|
|
let sum_per_sample = squared.sum(Some(1))?;
|
|
sum_per_sample.mean(&[], false)?.to_scalar::<f32>()?
|
|
}
|
|
DistanceMetric::Cosine => {
|
|
// 1 - cosine similarity
|
|
let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
|
|
let data_norm = data.pow(&two)?.sum(Some(1))?.sqrt()?;
|
|
let recon_norm = reconstructed.pow(&two)?.sum(Some(1))?.sqrt()?;
|
|
let dot_product = (data.clone() * reconstructed.clone())?.sum(Some(1))?;
|
|
let cosine_sim = (dot_product / (data_norm * recon_norm)?)?;
|
|
let one_minus_cosine = cosine_sim.mul_scalar(-1.0)?.add_scalar(1.0)?;
|
|
one_minus_cosine.mean(&[], false)?.to_scalar::<f32>()?
|
|
}
|
|
DistanceMetric::Manhattan => {
|
|
let diff = (data.clone() - reconstructed.clone())?;
|
|
let abs_diff = diff.abs()?;
|
|
let sum_per_sample = abs_diff.sum(Some(1))?;
|
|
sum_per_sample.mean(&[], false)?.to_scalar::<f32>()?
|
|
}
|
|
};
|
|
|
|
results.push((distance_metric.clone(), reconstruction_error));
|
|
|
|
println!(
|
|
"Distance metric {:?}: reconstruction error = {:.6}",
|
|
distance_metric, reconstruction_error
|
|
);
|
|
}
|
|
|
|
// All metrics should produce reasonable results
|
|
for (_, error) in &results {
|
|
assert!(
|
|
*error > 0.0 && *error < 100.0,
|
|
"Reconstruction error should be reasonable"
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal device randn issue"]
|
|
fn test_vq_online_updates() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Initial training data
|
|
let initial_data = Tensor::randn(&[1000, 32], &device)?;
|
|
|
|
let config = VQConfig {
|
|
codebook_size: 64,
|
|
vector_dim: 32,
|
|
max_iterations: 30,
|
|
tolerance: 1e-4,
|
|
initialization: CodebookInitialization::KMeansPlusPlus,
|
|
distance_metric: DistanceMetric::Euclidean,
|
|
};
|
|
|
|
let mut vq = VectorQuantizer::new(config);
|
|
vq.fit(&initial_data)?;
|
|
|
|
let initial_codebook = vq.codebook().clone();
|
|
|
|
// Enable online updates
|
|
vq.enable_online_updates(true, 0.01); // learning rate = 0.01
|
|
|
|
// Stream new data and update codebook online
|
|
for batch in 0..10 {
|
|
let new_data = Tensor::randn(&[100, 32], &device)?;
|
|
vq.update_online(&new_data)?;
|
|
}
|
|
|
|
let updated_codebook = vq.codebook().clone();
|
|
|
|
// Verify codebook has been updated
|
|
let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
|
|
let codebook_change = (initial_codebook - updated_codebook)?
|
|
.pow(&two)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
|
|
assert!(
|
|
codebook_change > 1e-4,
|
|
"Codebook should change with online updates, change = {}",
|
|
codebook_change
|
|
);
|
|
|
|
// Test encoding with updated codebook
|
|
let test_data = Tensor::randn(&[50, 32], &device)?;
|
|
let codes = vq.encode(&test_data)?;
|
|
let reconstructed = vq.decode(&codes)?;
|
|
|
|
let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
|
|
let mse = (test_data.clone() - reconstructed)?
|
|
.pow(&two)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
assert!(
|
|
mse < 1.0,
|
|
"Updated VQ should still provide reasonable reconstruction"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal device randn issue"]
|
|
fn test_vq_hierarchical_quantization() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
let data = Tensor::randn(&[2000, 256], &device)?;
|
|
|
|
// Create two-level hierarchical VQ
|
|
let coarse_config = VQConfig {
|
|
codebook_size: 64, // Coarse level
|
|
vector_dim: 256,
|
|
max_iterations: 50,
|
|
tolerance: 1e-5,
|
|
initialization: CodebookInitialization::KMeansPlusPlus,
|
|
distance_metric: DistanceMetric::Euclidean,
|
|
};
|
|
|
|
let fine_config = VQConfig {
|
|
codebook_size: 256, // Fine level
|
|
vector_dim: 256,
|
|
max_iterations: 50,
|
|
tolerance: 1e-5,
|
|
initialization: CodebookInitialization::KMeansPlusPlus,
|
|
distance_metric: DistanceMetric::Euclidean,
|
|
};
|
|
|
|
let mut coarse_vq = VectorQuantizer::new(coarse_config);
|
|
let mut fine_vq = VectorQuantizer::new(fine_config);
|
|
|
|
// First level: coarse quantization
|
|
coarse_vq.fit(&data)?;
|
|
let coarse_codes = coarse_vq.encode(&data)?;
|
|
let coarse_reconstructed = coarse_vq.decode(&coarse_codes)?;
|
|
|
|
// Second level: quantize residuals
|
|
let residuals = (data.clone() - coarse_reconstructed.clone())?;
|
|
fine_vq.fit(&residuals)?;
|
|
let fine_codes = fine_vq.encode(&residuals)?;
|
|
let fine_reconstructed = fine_vq.decode(&fine_codes)?;
|
|
|
|
// Final reconstruction
|
|
let final_reconstructed = (coarse_reconstructed.clone() + fine_reconstructed)?;
|
|
|
|
// Verify hierarchical VQ improves reconstruction
|
|
let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
|
|
let coarse_mse = (data.clone() - coarse_reconstructed.clone())?
|
|
.pow(&two)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
let hierarchical_mse = (data.clone() - final_reconstructed)?
|
|
.pow(&two)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
|
|
assert!(
|
|
hierarchical_mse < coarse_mse * 0.8,
|
|
"Hierarchical VQ should improve reconstruction: coarse={:.6}, hierarchical={:.6}",
|
|
coarse_mse,
|
|
hierarchical_mse
|
|
);
|
|
|
|
// Verify code storage efficiency
|
|
let total_codes_per_vector = 2; // One coarse + one fine code
|
|
let bits_per_code = (64_f32.log2().ceil() as usize) + (256_f32.log2().ceil() as usize);
|
|
let compression_ratio = (256 * 32) as f32 / bits_per_code as f32; // f32 = 32 bits
|
|
|
|
println!(
|
|
"Hierarchical VQ compression ratio: {:.2}x",
|
|
compression_ratio
|
|
);
|
|
assert!(
|
|
compression_ratio > 8.0,
|
|
"Should achieve significant compression"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal device randn issue"]
|
|
fn test_vq_codebook_pruning() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Create data with clusters of different sizes
|
|
let mut data_parts: Vec<Tensor> = vec![];
|
|
|
|
// Large cluster
|
|
let cluster1 = (Tensor::randn(&[800, 64], &device)?
|
|
+ Tensor::from_slice(&[2.0], &[1, 1], &device)?.broadcast_to(&[800, 64])?)?;
|
|
data_parts.push(cluster1);
|
|
|
|
// Medium cluster
|
|
let cluster2 = (Tensor::randn(&[150, 64], &device)?
|
|
+ Tensor::from_slice(&[-2.0], &[1, 1], &device)?.broadcast_to(&[150, 64])?)?;
|
|
data_parts.push(cluster2);
|
|
|
|
// Small cluster (should be pruned)
|
|
let cluster3 = (Tensor::randn(&[50, 64], &device)?
|
|
+ Tensor::from_slice(&[0.0, 4.0], &[1, 2], &device)?.broadcast_to(&[50, 64])?)?;
|
|
data_parts.push(cluster3);
|
|
|
|
let data = Tensor::cat(&data_parts, 0)?;
|
|
|
|
let config = VQConfig {
|
|
codebook_size: 128,
|
|
vector_dim: 64,
|
|
max_iterations: 100,
|
|
tolerance: 1e-6,
|
|
initialization: CodebookInitialization::KMeansPlusPlus,
|
|
distance_metric: DistanceMetric::Euclidean,
|
|
};
|
|
|
|
let mut vq = VectorQuantizer::new(config);
|
|
vq.fit(&data)?;
|
|
|
|
// Analyze codebook usage
|
|
let codes = vq.encode(&data)?;
|
|
let usage_stats = vq.analyze_codebook_usage(&codes)?;
|
|
|
|
assert_eq!(usage_stats.len(), 128);
|
|
|
|
// Find underutilized codewords (< 1% of data)
|
|
let min_usage_threshold = 0.01;
|
|
let underutilized: Vec<_> = usage_stats
|
|
.iter()
|
|
.enumerate()
|
|
.filter(|(_, usage)| **usage < min_usage_threshold)
|
|
.collect();
|
|
|
|
println!("Underutilized codewords: {}", underutilized.len());
|
|
|
|
// Prune underutilized codewords
|
|
let pruned_vq = vq.prune_codebook(min_usage_threshold)?;
|
|
|
|
assert!(
|
|
pruned_vq.codebook_size() < vq.codebook_size(),
|
|
"Pruned codebook should be smaller"
|
|
);
|
|
|
|
// Test that pruned VQ still works
|
|
let pruned_codes = pruned_vq.encode(&data)?;
|
|
let pruned_reconstructed = pruned_vq.decode(&pruned_codes)?;
|
|
|
|
let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
|
|
let original_mse = (data.clone() - vq.decode(&codes)?)?
|
|
.pow(&two)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
let pruned_mse = (data.clone() - pruned_reconstructed)?
|
|
.pow(&two)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
|
|
// Pruned VQ should have similar quality (within 20%)
|
|
assert!(
|
|
pruned_mse < original_mse * 1.2,
|
|
"Pruned VQ quality should be similar: original={:.6}, pruned={:.6}",
|
|
original_mse,
|
|
pruned_mse
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal device randn issue"]
|
|
fn test_vq_adaptive_codebook_size() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Test data with varying complexity
|
|
let simple_data = (Tensor::randn(&[1000, 32], &device)? * 0.1)?; // Low variance
|
|
let complex_data = (Tensor::randn(&[1000, 32], &device)? * 2.0)?; // High variance
|
|
|
|
let config = VQConfig {
|
|
codebook_size: 0, // Will be determined adaptively
|
|
vector_dim: 32,
|
|
max_iterations: 50,
|
|
tolerance: 1e-5,
|
|
initialization: CodebookInitialization::KMeansPlusPlus,
|
|
distance_metric: DistanceMetric::Euclidean,
|
|
};
|
|
|
|
let mut vq_simple = VectorQuantizer::new(config.clone());
|
|
let mut vq_complex = VectorQuantizer::new(config);
|
|
|
|
// Enable adaptive sizing
|
|
let target_distortion = 0.01;
|
|
vq_simple.enable_adaptive_sizing(target_distortion, 512)?; // max 512 codewords
|
|
vq_complex.enable_adaptive_sizing(target_distortion, 512)?;
|
|
|
|
// Fit with adaptive sizing
|
|
vq_simple.fit(&simple_data)?;
|
|
vq_complex.fit(&complex_data)?;
|
|
|
|
println!("Simple data codebook size: {}", vq_simple.codebook_size());
|
|
println!("Complex data codebook size: {}", vq_complex.codebook_size());
|
|
|
|
// Complex data should need more codewords
|
|
assert!(
|
|
vq_complex.codebook_size() > vq_simple.codebook_size(),
|
|
"Complex data should require larger codebook"
|
|
);
|
|
|
|
// Both should meet distortion target
|
|
let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
|
|
let simple_codes = vq_simple.encode(&simple_data)?;
|
|
let simple_reconstructed = vq_simple.decode(&simple_codes)?;
|
|
let simple_mse = (simple_data.clone() - simple_reconstructed)?
|
|
.pow(&two)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
|
|
let complex_codes = vq_complex.encode(&complex_data)?;
|
|
let complex_reconstructed = vq_complex.decode(&complex_codes)?;
|
|
let complex_mse = (complex_data.clone() - complex_reconstructed)?
|
|
.pow(&two)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
|
|
assert!(
|
|
simple_mse <= (target_distortion * 1.1) as f32,
|
|
"Simple VQ should meet distortion target: {:.6}",
|
|
simple_mse
|
|
);
|
|
assert!(
|
|
complex_mse <= (target_distortion * 1.1) as f32,
|
|
"Complex VQ should meet distortion target: {:.6}",
|
|
complex_mse
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal device randn issue"]
|
|
fn test_vq_batch_operations() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
let config = VQConfig {
|
|
codebook_size: 128,
|
|
vector_dim: 64,
|
|
max_iterations: 30,
|
|
tolerance: 1e-4,
|
|
initialization: CodebookInitialization::KMeansPlusPlus,
|
|
distance_metric: DistanceMetric::Euclidean,
|
|
};
|
|
|
|
let mut vq = VectorQuantizer::new(config);
|
|
|
|
// Train on batch
|
|
let train_data = Tensor::randn(&[2000, 64], &device)?;
|
|
vq.fit(&train_data)?;
|
|
|
|
// Test batch encoding with different shapes
|
|
let test_cases = vec![
|
|
Tensor::randn(&[100, 64], &device)?, // 2D batch
|
|
Tensor::randn(&[10, 10, 64], &device)?, // 3D batch
|
|
Tensor::randn(&[5, 4, 5, 64], &device)?, // 4D batch
|
|
];
|
|
|
|
for (i, test_data) in test_cases.iter().enumerate() {
|
|
let original_shape = test_data.shape().to_vec();
|
|
|
|
// Encode
|
|
let codes = vq.encode(test_data)?;
|
|
|
|
// Verify code shape (last dimension should be removed)
|
|
let expected_code_shape = original_shape[..original_shape.len() - 1].to_vec();
|
|
assert_eq!(
|
|
codes.shape(),
|
|
expected_code_shape.as_slice(),
|
|
"Case {}: code shape mismatch",
|
|
i
|
|
);
|
|
|
|
// Decode
|
|
let reconstructed = vq.decode(&codes)?;
|
|
|
|
// Verify reconstruction shape matches original
|
|
assert_eq!(
|
|
reconstructed.shape(),
|
|
test_data.shape(),
|
|
"Case {}: reconstruction shape mismatch",
|
|
i
|
|
);
|
|
|
|
// Verify reasonable reconstruction quality
|
|
let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
|
|
let mse = (test_data.clone() - reconstructed)?
|
|
.pow(&two)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
assert!(
|
|
mse < 2.0,
|
|
"Case {}: reconstruction quality poor: MSE={:.6}",
|
|
i,
|
|
mse
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal device randn issue"]
|
|
fn test_vq_serialization() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
let data = Tensor::randn(&[500, 128], &device)?;
|
|
|
|
let config = VQConfig {
|
|
codebook_size: 256,
|
|
vector_dim: 128,
|
|
max_iterations: 50,
|
|
tolerance: 1e-5,
|
|
initialization: CodebookInitialization::KMeansPlusPlus,
|
|
distance_metric: DistanceMetric::Euclidean,
|
|
};
|
|
|
|
let mut vq = VectorQuantizer::new(config);
|
|
vq.fit(&data)?;
|
|
|
|
// Test serialization
|
|
let serialized = vq.serialize()?;
|
|
|
|
// Deserialize into new VQ
|
|
let vq2 = VectorQuantizer::deserialize(&serialized)?;
|
|
|
|
// Verify both VQs produce identical results
|
|
let codes1 = vq.encode(&data)?;
|
|
let codes2 = vq2.encode(&data)?;
|
|
|
|
let code_diff = (codes1.to_dtype(DType::F32)? - codes2.to_dtype(DType::F32)?)?
|
|
.abs()?
|
|
.max()?
|
|
.to_scalar::<f32>()?;
|
|
|
|
assert!(
|
|
code_diff < 1e-6,
|
|
"Serialized VQ should produce identical codes"
|
|
);
|
|
|
|
// Verify codebooks are identical
|
|
let codebook_diff = (vq.codebook().clone() - vq2.codebook().clone())?
|
|
.abs()?
|
|
.max()?
|
|
.to_scalar::<f32>()?;
|
|
|
|
assert!(
|
|
codebook_diff < 1e-6,
|
|
"Serialized codebook should be identical"
|
|
);
|
|
|
|
Ok(())
|
|
}
|