106 lines
3.3 KiB
Rust
106 lines
3.3 KiB
Rust
// Basic compression tests - simplified to match current API
|
|
use rtx_compress::{
|
|
Result,
|
|
pipeline::{
|
|
CompressionPipeline, CompressionPipelineConfig, CompressionStrategy, HardwareTarget,
|
|
},
|
|
quantization::{
|
|
CalibrationConfig, PostTrainingQuantizer, QuantizationConfig, QuantizationScheme,
|
|
},
|
|
};
|
|
use rtx_tensor::{Device, Tensor};
|
|
use std::collections::HashMap;
|
|
|
|
#[test]
|
|
fn test_compression_pipeline_creation() {
|
|
let config = CompressionPipelineConfig {
|
|
strategy: CompressionStrategy::Balanced,
|
|
target_compression_ratio: 4.0,
|
|
target_accuracy_retention: 0.95,
|
|
progressive: false,
|
|
num_stages: 2,
|
|
validation_size: 100,
|
|
hardware_target: HardwareTarget::CPU,
|
|
};
|
|
let pipeline = CompressionPipeline::new(config);
|
|
assert!(pipeline.is_ok(), "Should create compression pipeline");
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal device randn issue"]
|
|
fn test_post_training_quantization() {
|
|
let device = Device::try_default().unwrap();
|
|
|
|
let calibration = CalibrationConfig::new(100, 0.99);
|
|
let config = QuantizationConfig::new(QuantizationScheme::INT8, calibration);
|
|
|
|
let quantizer = PostTrainingQuantizer::new(config);
|
|
assert!(quantizer.is_ok(), "Should create post-training quantizer");
|
|
|
|
let mut quantizer = quantizer.unwrap();
|
|
|
|
// Create calibration data
|
|
let calibration_data: Vec<Tensor> = (0..10)
|
|
.map(|_| Tensor::randn(&[64, 64], &device).unwrap())
|
|
.collect();
|
|
|
|
let calibrate_result = quantizer.calibrate(&calibration_data);
|
|
assert!(calibrate_result.is_ok(), "Should calibrate quantizer");
|
|
|
|
// Quantize a tensor
|
|
let tensor = Tensor::randn(&[64, 64], &device).unwrap();
|
|
let quantized = quantizer.quantize(&tensor);
|
|
assert!(quantized.is_ok(), "Should quantize tensor");
|
|
|
|
let quantized_tensor = quantized.unwrap();
|
|
|
|
// Dequantize
|
|
let dequantized = quantizer.dequantize(&quantized_tensor);
|
|
assert!(dequantized.is_ok(), "Should dequantize tensor");
|
|
|
|
// Check shape preservation
|
|
let deq = dequantized.unwrap();
|
|
assert_eq!(deq.shape().dims(), tensor.shape().dims());
|
|
}
|
|
|
|
#[test]
|
|
fn test_compression_pipeline_auto() {
|
|
let pipeline = CompressionPipeline::auto(4.0, 0.95, HardwareTarget::GPU);
|
|
assert!(pipeline.is_ok(), "Should create auto-tuned pipeline");
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal device randn issue"]
|
|
fn test_end_to_end_compression_pipeline() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Create a simple model
|
|
let mut model = HashMap::new();
|
|
model.insert("layer1".to_string(), Tensor::randn(&[512, 512], &device)?);
|
|
model.insert("layer2".to_string(), Tensor::randn(&[512, 256], &device)?);
|
|
model.insert("layer3".to_string(), Tensor::randn(&[256, 128], &device)?);
|
|
|
|
let config = CompressionPipelineConfig {
|
|
strategy: CompressionStrategy::Size,
|
|
target_compression_ratio: 4.0,
|
|
target_accuracy_retention: 0.90,
|
|
progressive: false,
|
|
num_stages: 1,
|
|
validation_size: 100,
|
|
hardware_target: HardwareTarget::CPU,
|
|
};
|
|
|
|
let mut pipeline = CompressionPipeline::new(config)?;
|
|
|
|
// Apply compression
|
|
let result = pipeline.compress(&model, None, None)?;
|
|
|
|
// Verify compression was achieved
|
|
assert!(
|
|
result.statistics.compression_ratio >= 1.0,
|
|
"Should achieve some compression"
|
|
);
|
|
|
|
Ok(())
|
|
}
|