307 lines
9.3 KiB
Rust
307 lines
9.3 KiB
Rust
//! Tests for model pruning functionality
|
|
//!
|
|
//! Following strict TDD - tests define expected behavior before implementation
|
|
#![cfg(feature = "disabled_tests")]
|
|
|
|
use approx::assert_abs_diff_eq;
|
|
use rtx_compress::{
|
|
Result,
|
|
pruning::{
|
|
MagnitudePruner, PruningConfig, PruningGranularity, PruningMethod, StructuredPruner,
|
|
},
|
|
};
|
|
use rtx_tensor::{Device, Tensor};
|
|
|
|
#[test]
|
|
fn test_magnitude_pruner_creation() {
|
|
let device = Device::cpu();
|
|
|
|
let config = PruningConfig {
|
|
sparsity: 0.5, // 50% sparsity
|
|
structured: false,
|
|
granularity: PruningGranularity::Unstructured,
|
|
preserve_gradients: true,
|
|
};
|
|
|
|
let pruner = MagnitudePruner::new(config, &device);
|
|
assert!(pruner.is_ok());
|
|
|
|
let pruner = pruner.unwrap();
|
|
assert_eq!(pruner.sparsity(), 0.5);
|
|
assert!(!pruner.is_structured());
|
|
}
|
|
|
|
#[test]
|
|
fn test_unstructured_magnitude_pruning() {
|
|
let device = Device::cpu();
|
|
|
|
// Create a weight tensor with known values
|
|
let weights_data = vec![
|
|
0.1, -0.5, 0.3, -0.2, // Small and large magnitudes
|
|
0.9, -0.05, 0.7, -0.8, // Mixed values
|
|
0.01, -0.02, 0.6, -0.4, // Very small and medium
|
|
0.15, -0.95, 0.03, -0.35,
|
|
];
|
|
let weights = Tensor::from_data(weights_data.clone(), &[4, 4], &device).unwrap();
|
|
|
|
// Prune 50% of weights (8 out of 16)
|
|
let config = PruningConfig {
|
|
sparsity: 0.5,
|
|
structured: false,
|
|
granularity: PruningGranularity::Unstructured,
|
|
preserve_gradients: false,
|
|
};
|
|
|
|
let pruner = MagnitudePruner::new(config, &device).unwrap();
|
|
let mask = pruner.compute_mask(&weights).unwrap();
|
|
|
|
// Apply mask
|
|
let pruned = pruner.apply_mask(&weights, &mask).unwrap();
|
|
|
|
// Check that exactly 50% of weights are zero
|
|
let pruned_data = pruned.to_vec().unwrap();
|
|
let num_zeros = pruned_data.iter().filter(|&&x| x == 0.0).count();
|
|
assert_eq!(num_zeros, 8);
|
|
|
|
// Check that smallest magnitude weights were pruned
|
|
// Values with smallest magnitude: 0.01, -0.02, 0.03, -0.05, 0.1, 0.15, -0.2, 0.3
|
|
// Should keep: -0.5, 0.9, 0.7, -0.8, 0.6, -0.4, -0.95, -0.35
|
|
for (i, val) in pruned_data.iter().enumerate() {
|
|
let original = weights_data[i];
|
|
if original.abs() < 0.35 {
|
|
assert_eq!(
|
|
*val, 0.0,
|
|
"Small magnitude weight at index {} should be pruned",
|
|
i
|
|
);
|
|
} else {
|
|
assert_abs_diff_eq!(*val, original, epsilon = 1e-6);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_structured_pruning_channels() {
|
|
let device = Device::cpu();
|
|
|
|
// Create a Conv2d weight tensor [out_channels, in_channels, kernel_h, kernel_w]
|
|
let weights = Tensor::randn(&[8, 3, 3, 3], &device).unwrap();
|
|
|
|
// Prune 50% of output channels (4 out of 8)
|
|
let config = PruningConfig {
|
|
sparsity: 0.5,
|
|
structured: true,
|
|
granularity: PruningGranularity::Channel,
|
|
preserve_gradients: false,
|
|
};
|
|
|
|
let pruner = StructuredPruner::new(config, &device).unwrap();
|
|
let mask = pruner.compute_mask(&weights).unwrap();
|
|
let pruned = pruner.apply_mask(&weights, &mask).unwrap();
|
|
|
|
// Check that entire channels are pruned
|
|
let pruned_shape = pruned.shape().dims();
|
|
assert_eq!(pruned_shape, &[8, 3, 3, 3]);
|
|
|
|
// Count pruned channels
|
|
let pruned_data = pruned.to_vec().unwrap();
|
|
let channel_size = 3 * 3 * 3;
|
|
let mut pruned_channels = 0;
|
|
|
|
for ch in 0..8 {
|
|
let start = ch * channel_size;
|
|
let end = start + channel_size;
|
|
let channel_sum: f32 = pruned_data[start..end].iter().map(|x| x.abs()).sum();
|
|
if channel_sum == 0.0 {
|
|
pruned_channels += 1;
|
|
}
|
|
}
|
|
|
|
assert_eq!(pruned_channels, 4, "Should prune exactly 4 channels");
|
|
}
|
|
|
|
#[test]
|
|
fn test_structured_pruning_filters() {
|
|
let device = Device::cpu();
|
|
|
|
// Create weight tensor for structured filter pruning
|
|
let weights = Tensor::randn(&[64, 64, 3, 3], &device).unwrap();
|
|
|
|
let config = PruningConfig {
|
|
sparsity: 0.25, // Prune 25% of filters
|
|
structured: true,
|
|
granularity: PruningGranularity::Filter,
|
|
preserve_gradients: false,
|
|
};
|
|
|
|
let pruner = StructuredPruner::new(config, &device).unwrap();
|
|
let indices = pruner.get_pruned_indices(&weights).unwrap();
|
|
|
|
// Should identify 16 filters to prune (25% of 64)
|
|
assert_eq!(indices.len(), 16);
|
|
|
|
// Apply pruning
|
|
let pruned = pruner.prune(&weights).unwrap();
|
|
|
|
// Verify filters are zeroed
|
|
let pruned_data = pruned.to_vec().unwrap();
|
|
let filter_size = 64 * 3 * 3;
|
|
|
|
for &idx in indices.iter() {
|
|
let start = idx * filter_size;
|
|
let end = start + filter_size;
|
|
let filter_sum: f32 = pruned_data[start..end].iter().map(|x| x.abs()).sum();
|
|
assert_eq!(
|
|
filter_sum, 0.0,
|
|
"Filter {} should be completely pruned",
|
|
idx
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_gradual_magnitude_pruning() {
|
|
let device = Device::cpu();
|
|
|
|
// Test gradual pruning over multiple iterations
|
|
let weights = Tensor::randn(&[100, 100], &device).unwrap();
|
|
|
|
let mut pruner = MagnitudePruner::with_schedule(
|
|
0.9, // Final sparsity
|
|
100, // Number of steps
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Initial pruning should be minimal
|
|
let mask_initial = pruner.compute_mask_at_step(&weights, 0).unwrap();
|
|
let sparsity_initial = pruner.get_current_sparsity(0);
|
|
assert!(sparsity_initial < 0.1, "Initial sparsity should be low");
|
|
|
|
// Middle of training
|
|
let mask_middle = pruner.compute_mask_at_step(&weights, 50).unwrap();
|
|
let sparsity_middle = pruner.get_current_sparsity(50);
|
|
assert!(
|
|
sparsity_middle > 0.3 && sparsity_middle < 0.6,
|
|
"Middle sparsity should be moderate"
|
|
);
|
|
|
|
// End of training
|
|
let mask_final = pruner.compute_mask_at_step(&weights, 99).unwrap();
|
|
let sparsity_final = pruner.get_current_sparsity(99);
|
|
assert_abs_diff_eq!(sparsity_final, 0.9, epsilon = 0.01);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pruning_with_importance_scores() {
|
|
let device = Device::cpu();
|
|
|
|
// Create weights and importance scores
|
|
let weights = Tensor::randn(&[10, 10], &device).unwrap();
|
|
let importance = Tensor::rand(&[10, 10], &device).unwrap(); // Random importance scores
|
|
|
|
let config = PruningConfig {
|
|
sparsity: 0.3,
|
|
structured: false,
|
|
granularity: PruningGranularity::Unstructured,
|
|
preserve_gradients: false,
|
|
};
|
|
|
|
let pruner = MagnitudePruner::new(config, &device).unwrap();
|
|
|
|
// Prune based on importance scores instead of magnitude
|
|
let mask = pruner
|
|
.compute_mask_with_importance(&weights, &importance)
|
|
.unwrap();
|
|
let pruned = pruner.apply_mask(&weights, &mask).unwrap();
|
|
|
|
// Check sparsity
|
|
let pruned_data = pruned.to_vec().unwrap();
|
|
let num_zeros = pruned_data.iter().filter(|&&x| x == 0.0).count();
|
|
let expected_zeros = (100.0 * 0.3) as usize;
|
|
assert_eq!(num_zeros, expected_zeros);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pruning_mask_persistence() {
|
|
let device = Device::cpu();
|
|
|
|
let config = PruningConfig {
|
|
sparsity: 0.4,
|
|
structured: false,
|
|
granularity: PruningGranularity::Unstructured,
|
|
preserve_gradients: true,
|
|
};
|
|
|
|
let pruner = MagnitudePruner::new(config, &device).unwrap();
|
|
let weights = Tensor::randn(&[50, 50], &device).unwrap();
|
|
|
|
// Compute and save mask
|
|
let mask = pruner.compute_mask(&weights).unwrap();
|
|
let mask_data = mask.to_vec().unwrap();
|
|
|
|
// Apply mask multiple times - should be consistent
|
|
let pruned1 = pruner.apply_mask(&weights, &mask).unwrap();
|
|
let pruned2 = pruner.apply_mask(&weights, &mask).unwrap();
|
|
|
|
let data1 = pruned1.to_vec().unwrap();
|
|
let data2 = pruned2.to_vec().unwrap();
|
|
|
|
for i in 0..data1.len() {
|
|
assert_abs_diff_eq!(data1[i], data2[i], epsilon = 1e-6);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_fine_grained_pruning() {
|
|
let device = Device::cpu();
|
|
|
|
// Test N:M sparsity pattern (e.g., 2:4 for NVIDIA Ampere)
|
|
let weights = Tensor::randn(&[64, 64], &device).unwrap();
|
|
|
|
let pruner = MagnitudePruner::with_nm_sparsity(
|
|
2, // N (number of non-zeros)
|
|
4, // M (block size)
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
let mask = pruner.compute_mask(&weights).unwrap();
|
|
let pruned = pruner.apply_mask(&weights, &mask).unwrap();
|
|
|
|
// Verify 2:4 pattern
|
|
let pruned_data = pruned.to_vec().unwrap();
|
|
|
|
// Check each block of 4 elements has exactly 2 zeros
|
|
for i in (0..pruned_data.len()).step_by(4) {
|
|
if i + 3 < pruned_data.len() {
|
|
let block = &pruned_data[i..i + 4];
|
|
let zeros = block.iter().filter(|&&x| x == 0.0).count();
|
|
assert_eq!(zeros, 2, "Each block of 4 should have exactly 2 zeros");
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_pruning_statistics() {
|
|
let device = Device::cpu();
|
|
|
|
let config = PruningConfig {
|
|
sparsity: 0.6,
|
|
structured: false,
|
|
granularity: PruningGranularity::Unstructured,
|
|
preserve_gradients: false,
|
|
};
|
|
|
|
let pruner = MagnitudePruner::new(config, &device).unwrap();
|
|
let weights = Tensor::randn(&[128, 128], &device).unwrap();
|
|
|
|
let stats = pruner.analyze(&weights).unwrap();
|
|
|
|
// Check statistics
|
|
assert_eq!(stats.total_parameters, 128 * 128);
|
|
assert_eq!(stats.target_sparsity, 0.6);
|
|
assert!(stats.achieved_sparsity >= 0.0 && stats.achieved_sparsity <= 1.0);
|
|
assert!(stats.compression_ratio > 1.0); // Should achieve some compression
|
|
}
|