743 lines
24 KiB
Rust
743 lines
24 KiB
Rust
//! Comprehensive tests for magnitude-based unstructured pruning
|
|
//!
|
|
//! Tests cover:
|
|
//! - Global and layer-wise magnitude pruning
|
|
//! - Gradual pruning with different schedules
|
|
#![cfg(feature = "disabled_tests")]
|
|
//! - Top-k sparsity patterns
|
|
//! - Gradient-based importance (SNIP)
|
|
//! - Random and structured-random pruning
|
|
//! - Fine-tuning with sparse gradients
|
|
//! - Sparse-to-sparse training continuation
|
|
|
|
use rtx_autograd::{no_grad, tensor_with_grad};
|
|
use rtx_compress::{
|
|
Result,
|
|
pruning::{
|
|
GradualPruningSchedule, PruningCriterion, PruningMask, SparsityPattern, UnstructuredPruner,
|
|
UnstructuredPruningConfig,
|
|
},
|
|
};
|
|
use rtx_tensor::{Device, Tensor};
|
|
use std::collections::HashMap;
|
|
|
|
#[cfg(test)]
|
|
mod unstructured_pruning_tests {
|
|
use super::*;
|
|
|
|
fn create_test_tensor(shape: &[usize], seed: u64) -> Result<Tensor> {
|
|
let device = Device::try_default()?;
|
|
let mut data = Vec::new();
|
|
let total_elements: usize = shape.iter().product();
|
|
|
|
// Create deterministic pseudo-random data
|
|
for i in 0..total_elements {
|
|
let val = (((i as u64).wrapping_mul(seed) % 1000) as f32 / 1000.0 - 0.5) * 2.0;
|
|
data.push(val);
|
|
}
|
|
|
|
Tensor::from_slice(&data, shape, &device)
|
|
}
|
|
|
|
fn create_test_model(seed: u64) -> Result<HashMap<String, Tensor>> {
|
|
let mut model = HashMap::new();
|
|
|
|
model.insert(
|
|
"conv1.weight".to_string(),
|
|
create_test_tensor(&[64, 3, 3, 3], seed)?,
|
|
);
|
|
model.insert(
|
|
"conv1.bias".to_string(),
|
|
create_test_tensor(&[64], seed + 1)?,
|
|
);
|
|
model.insert(
|
|
"conv2.weight".to_string(),
|
|
create_test_tensor(&[128, 64, 3, 3], seed + 2)?,
|
|
);
|
|
model.insert(
|
|
"conv2.bias".to_string(),
|
|
create_test_tensor(&[128], seed + 3)?,
|
|
);
|
|
model.insert(
|
|
"fc1.weight".to_string(),
|
|
create_test_tensor(&[256, 512], seed + 4)?,
|
|
);
|
|
model.insert(
|
|
"fc1.bias".to_string(),
|
|
create_test_tensor(&[256], seed + 5)?,
|
|
);
|
|
model.insert(
|
|
"fc2.weight".to_string(),
|
|
create_test_tensor(&[10, 256], seed + 6)?,
|
|
);
|
|
model.insert("fc2.bias".to_string(), create_test_tensor(&[10], seed + 7)?);
|
|
|
|
Ok(model)
|
|
}
|
|
|
|
fn calculate_sparsity(tensor: &Tensor) -> Result<f32> {
|
|
let values = tensor.to_vec::<f32>()?;
|
|
let zero_count = values.iter().filter(|&&x| x.abs() < 1e-8).count();
|
|
Ok(zero_count as f32 / values.len() as f32)
|
|
}
|
|
|
|
#[test]
|
|
fn test_global_magnitude_pruning() -> Result<()> {
|
|
let config = UnstructuredPruningConfig::new(
|
|
PruningCriterion::GlobalMagnitude,
|
|
0.8, // 80% sparsity
|
|
);
|
|
|
|
let pruner = UnstructuredPruner::new(config)?;
|
|
let model = create_test_model(42)?;
|
|
|
|
let pruned_model = pruner.prune_model(&model)?;
|
|
|
|
// Check that all tensors maintain their shape
|
|
for (name, original) in &model {
|
|
let pruned = &pruned_model[name];
|
|
assert_eq!(
|
|
original.shape(),
|
|
pruned.shape(),
|
|
"Shape mismatch for {}",
|
|
name
|
|
);
|
|
}
|
|
|
|
// Calculate overall sparsity
|
|
let mut total_elements = 0;
|
|
let mut total_zeros = 0;
|
|
|
|
for (_, tensor) in &pruned_model {
|
|
let values = tensor.to_vec::<f32>()?;
|
|
total_elements += values.len();
|
|
total_zeros += values.iter().filter(|&&x| x.abs() < 1e-8).count();
|
|
}
|
|
|
|
let overall_sparsity = total_zeros as f32 / total_elements as f32;
|
|
assert!(
|
|
(overall_sparsity - 0.8).abs() < 0.01,
|
|
"Expected ~80% sparsity, got {:.2}%",
|
|
overall_sparsity * 100.0
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_layer_wise_magnitude_pruning() -> Result<()> {
|
|
let mut layer_sparsities = HashMap::new();
|
|
layer_sparsities.insert("conv1.weight".to_string(), 0.5);
|
|
layer_sparsities.insert("conv2.weight".to_string(), 0.7);
|
|
layer_sparsities.insert("fc1.weight".to_string(), 0.9);
|
|
layer_sparsities.insert("fc2.weight".to_string(), 0.6);
|
|
|
|
let config = UnstructuredPruningConfig::new_layer_wise(
|
|
PruningCriterion::LocalMagnitude,
|
|
layer_sparsities,
|
|
);
|
|
|
|
let pruner = UnstructuredPruner::new(config)?;
|
|
let model = create_test_model(42)?;
|
|
|
|
let pruned_model = pruner.prune_model(&model)?;
|
|
|
|
// Check individual layer sparsities
|
|
let conv1_sparsity = calculate_sparsity(&pruned_model["conv1.weight"])?;
|
|
let conv2_sparsity = calculate_sparsity(&pruned_model["conv2.weight"])?;
|
|
let fc1_sparsity = calculate_sparsity(&pruned_model["fc1.weight"])?;
|
|
let fc2_sparsity = calculate_sparsity(&pruned_model["fc2.weight"])?;
|
|
|
|
assert!(
|
|
(conv1_sparsity - 0.5).abs() < 0.02,
|
|
"Conv1 sparsity: {:.3}",
|
|
conv1_sparsity
|
|
);
|
|
assert!(
|
|
(conv2_sparsity - 0.7).abs() < 0.02,
|
|
"Conv2 sparsity: {:.3}",
|
|
conv2_sparsity
|
|
);
|
|
assert!(
|
|
(fc1_sparsity - 0.9).abs() < 0.02,
|
|
"FC1 sparsity: {:.3}",
|
|
fc1_sparsity
|
|
);
|
|
assert!(
|
|
(fc2_sparsity - 0.6).abs() < 0.02,
|
|
"FC2 sparsity: {:.3}",
|
|
fc2_sparsity
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_gradual_pruning_schedule() -> Result<()> {
|
|
let schedule = GradualPruningSchedule::polynomial(
|
|
0.0, // Start with no pruning
|
|
0.9, // End with 90% sparsity
|
|
1000, // Over 1000 steps
|
|
3.0, // Cubic schedule
|
|
);
|
|
|
|
let config = UnstructuredPruningConfig::new_with_schedule(
|
|
PruningCriterion::GlobalMagnitude,
|
|
schedule,
|
|
);
|
|
|
|
let mut pruner = UnstructuredPruner::new(config)?;
|
|
let tensor = create_test_tensor(&[128, 256], 42)?;
|
|
|
|
// Test gradual pruning over multiple steps
|
|
let mut sparsities = Vec::new();
|
|
let steps = [0, 100, 300, 600, 999];
|
|
|
|
for &step in &steps {
|
|
pruner.set_current_step(step);
|
|
let pruned = pruner.prune_tensor(&tensor, "test")?;
|
|
let sparsity = calculate_sparsity(&pruned)?;
|
|
sparsities.push(sparsity);
|
|
}
|
|
|
|
// Sparsity should increase monotonically
|
|
for i in 1..sparsities.len() {
|
|
assert!(
|
|
sparsities[i] >= sparsities[i - 1],
|
|
"Sparsity should increase: step {} = {:.3}, step {} = {:.3}",
|
|
steps[i - 1],
|
|
sparsities[i - 1],
|
|
steps[i],
|
|
sparsities[i]
|
|
);
|
|
}
|
|
|
|
// Final sparsity should be close to target
|
|
assert!(
|
|
(sparsities.last().unwrap() - 0.9).abs() < 0.05,
|
|
"Final sparsity should be ~90%, got {:.1}%",
|
|
sparsities.last().unwrap() * 100.0
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_snip_gradient_based_pruning() -> Result<()> {
|
|
let config = UnstructuredPruningConfig::new(
|
|
PruningCriterion::SNIP, // Single-shot Network Pruning
|
|
0.8,
|
|
);
|
|
|
|
let mut pruner = UnstructuredPruner::new(config)?;
|
|
|
|
// Create tensor with gradients for SNIP
|
|
let tensor = create_test_tensor(&[64, 128], 42)?;
|
|
let gradient = create_test_tensor(&[64, 128], 123)?; // Simulate gradient
|
|
|
|
let pruned = pruner.prune_tensor_with_gradient(&tensor, "test", &gradient)?;
|
|
let sparsity = calculate_sparsity(&pruned)?;
|
|
|
|
assert!(
|
|
(sparsity - 0.8).abs() < 0.02,
|
|
"SNIP sparsity should be ~80%, got {:.1}%",
|
|
sparsity * 100.0
|
|
);
|
|
|
|
// Verify that SNIP preserves weights with high gradient * weight product
|
|
let tensor_vals = tensor.to_vec::<f32>()?;
|
|
let grad_vals = gradient.to_vec::<f32>()?;
|
|
let pruned_vals = pruned.to_vec::<f32>()?;
|
|
|
|
let mut preserved_high_importance = 0;
|
|
let mut total_high_importance = 0;
|
|
|
|
for ((w, g), p) in tensor_vals
|
|
.iter()
|
|
.zip(grad_vals.iter())
|
|
.zip(pruned_vals.iter())
|
|
{
|
|
let importance = w.abs() * g.abs();
|
|
if importance > 0.1 {
|
|
// Arbitrary threshold for "high importance"
|
|
total_high_importance += 1;
|
|
if p.abs() > 1e-8 {
|
|
preserved_high_importance += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
if total_high_importance > 0 {
|
|
let preservation_rate = preserved_high_importance as f32 / total_high_importance as f32;
|
|
assert!(
|
|
preservation_rate > 0.3,
|
|
"SNIP should preserve high-importance weights better: {:.1}%",
|
|
preservation_rate * 100.0
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_top_k_sparsity_pattern() -> Result<()> {
|
|
let config = UnstructuredPruningConfig::new_with_pattern(
|
|
PruningCriterion::GlobalMagnitude,
|
|
0.75, // 75% sparsity
|
|
SparsityPattern::TopK { k: 1000 }, // Keep top 1000 weights per tensor
|
|
);
|
|
|
|
let pruner = UnstructuredPruner::new(config)?;
|
|
let tensor = create_test_tensor(&[50, 100], 42)?; // 5000 elements total
|
|
|
|
let pruned = pruner.prune_tensor(&tensor, "test")?;
|
|
let values = pruned.to_vec::<f32>()?;
|
|
|
|
// Count non-zero values
|
|
let non_zero_count = values.iter().filter(|&&x| x.abs() > 1e-8).count();
|
|
|
|
// Should keep approximately k values (allowing for ties)
|
|
assert!(
|
|
non_zero_count >= 900 && non_zero_count <= 1100,
|
|
"Top-K should keep ~1000 values, got {}",
|
|
non_zero_count
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_random_pruning() -> Result<()> {
|
|
let config = UnstructuredPruningConfig::new(
|
|
PruningCriterion::Random { seed: 42 },
|
|
0.6, // 60% sparsity
|
|
);
|
|
|
|
let pruner = UnstructuredPruner::new(config)?;
|
|
let tensor = create_test_tensor(&[32, 64], 42)?;
|
|
|
|
let pruned1 = pruner.prune_tensor(&tensor, "test1")?;
|
|
let pruned2 = pruner.prune_tensor(&tensor, "test2")?;
|
|
|
|
// Both should have same sparsity
|
|
let sparsity1 = calculate_sparsity(&pruned1)?;
|
|
let sparsity2 = calculate_sparsity(&pruned2)?;
|
|
|
|
assert!((sparsity1 - 0.6).abs() < 0.05);
|
|
assert!((sparsity2 - 0.6).abs() < 0.05);
|
|
|
|
// Should produce identical results (deterministic)
|
|
let vals1 = pruned1.to_vec::<f32>()?;
|
|
let vals2 = pruned2.to_vec::<f32>()?;
|
|
|
|
let differences = vals1
|
|
.iter()
|
|
.zip(vals2.iter())
|
|
.filter(|(&a, &b)| (a - b).abs() > 1e-6)
|
|
.count();
|
|
|
|
assert_eq!(
|
|
differences, 0,
|
|
"Random pruning should be deterministic with same seed"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_structured_random_pruning() -> Result<()> {
|
|
let config = UnstructuredPruningConfig::new_with_pattern(
|
|
PruningCriterion::Random { seed: 42 },
|
|
0.75, // 75% sparsity
|
|
SparsityPattern::Block { block_size: 4 }, // 4x4 blocks
|
|
);
|
|
|
|
let pruner = UnstructuredPruner::new(config)?;
|
|
let tensor = create_test_tensor(&[32, 32], 42)?; // Square tensor for easy blocking
|
|
|
|
let pruned = pruner.prune_tensor(&tensor, "test")?;
|
|
let sparsity = calculate_sparsity(&pruned)?;
|
|
|
|
assert!(
|
|
(sparsity - 0.75).abs() < 0.1,
|
|
"Block sparsity: {:.2}",
|
|
sparsity
|
|
);
|
|
|
|
// Verify block structure (simplified check)
|
|
let values = pruned.to_vec::<f32>()?;
|
|
let mut block_zero_count = 0;
|
|
let mut total_blocks = 0;
|
|
|
|
// Check 4x4 blocks
|
|
for block_row in 0..8 {
|
|
// 32/4 = 8
|
|
for block_col in 0..8 {
|
|
let mut block_is_zero = true;
|
|
for i in 0..4 {
|
|
for j in 0..4 {
|
|
let row = block_row * 4 + i;
|
|
let col = block_col * 4 + j;
|
|
let idx = row * 32 + col;
|
|
if values[idx].abs() > 1e-8 {
|
|
block_is_zero = false;
|
|
break;
|
|
}
|
|
}
|
|
if !block_is_zero {
|
|
break;
|
|
}
|
|
}
|
|
if block_is_zero {
|
|
block_zero_count += 1;
|
|
}
|
|
total_blocks += 1;
|
|
}
|
|
}
|
|
|
|
let block_sparsity = block_zero_count as f32 / total_blocks as f32;
|
|
// Block pruning should create some fully zero blocks
|
|
assert!(
|
|
block_sparsity > 0.1,
|
|
"Should have some zero blocks: {:.2}",
|
|
block_sparsity
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_magnitude_vs_gradient_pruning() -> Result<()> {
|
|
let tensor = create_test_tensor(&[16, 32], 42)?;
|
|
let gradient = create_test_tensor(&[16, 32], 123)?;
|
|
|
|
// Magnitude-based pruning
|
|
let mag_config = UnstructuredPruningConfig::new(PruningCriterion::GlobalMagnitude, 0.5);
|
|
let mag_pruner = UnstructuredPruner::new(mag_config)?;
|
|
let mag_pruned = mag_pruner.prune_tensor(&tensor, "test")?;
|
|
|
|
// Gradient-based pruning (SNIP)
|
|
let snip_config = UnstructuredPruningConfig::new(PruningCriterion::SNIP, 0.5);
|
|
let mut snip_pruner = UnstructuredPruner::new(snip_config)?;
|
|
let snip_pruned = snip_pruner.prune_tensor_with_gradient(&tensor, "test", &gradient)?;
|
|
|
|
// Both should achieve similar sparsity
|
|
let mag_sparsity = calculate_sparsity(&mag_pruned)?;
|
|
let snip_sparsity = calculate_sparsity(&snip_pruned)?;
|
|
|
|
assert!((mag_sparsity - 0.5).abs() < 0.05);
|
|
assert!((snip_sparsity - 0.5).abs() < 0.05);
|
|
|
|
// But should preserve different weights
|
|
let mag_vals = mag_pruned.to_vec::<f32>()?;
|
|
let snip_vals = snip_pruned.to_vec::<f32>()?;
|
|
|
|
let different_preserved = mag_vals
|
|
.iter()
|
|
.zip(snip_vals.iter())
|
|
.filter(|(&m, &s)| (m.abs() > 1e-8) != (s.abs() > 1e-8))
|
|
.count();
|
|
|
|
assert!(
|
|
different_preserved > 100,
|
|
"Magnitude and SNIP should preserve different weights: {} differences",
|
|
different_preserved
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_pruning_mask_generation() -> Result<()> {
|
|
let config = UnstructuredPruningConfig::new_with_mask_generation(
|
|
PruningCriterion::LocalMagnitude,
|
|
0.7,
|
|
true, // Generate masks
|
|
);
|
|
|
|
let pruner = UnstructuredPruner::new(config)?;
|
|
let tensor = create_test_tensor(&[24, 48], 42)?;
|
|
|
|
let result = pruner.prune_tensor_with_mask(&tensor, "test")?;
|
|
|
|
// Check that mask is binary
|
|
let mask_values = result.mask.to_vec::<f32>()?;
|
|
for &val in &mask_values {
|
|
assert!(
|
|
val == 0.0 || val == 1.0,
|
|
"Mask should be binary, found {}",
|
|
val
|
|
);
|
|
}
|
|
|
|
// Check that pruned tensor equals original tensor masked
|
|
let original_values = tensor.to_vec::<f32>()?;
|
|
let pruned_values = result.pruned_tensor.to_vec::<f32>()?;
|
|
|
|
for ((orig, pruned), &mask) in original_values
|
|
.iter()
|
|
.zip(pruned_values.iter())
|
|
.zip(mask_values.iter())
|
|
{
|
|
if mask == 0.0 {
|
|
assert_eq!(*pruned, 0.0, "Masked weights should be zero");
|
|
} else {
|
|
assert_eq!(*pruned, *orig, "Unmasked weights should be preserved");
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_fine_tuning_mask_persistence() -> Result<()> {
|
|
let config = UnstructuredPruningConfig::new_with_mask_generation(
|
|
PruningCriterion::GlobalMagnitude,
|
|
0.8,
|
|
true,
|
|
);
|
|
|
|
let pruner = UnstructuredPruner::new(config)?;
|
|
let tensor = create_test_tensor(&[16, 32], 42)?;
|
|
|
|
let result = pruner.prune_tensor_with_mask(&tensor, "test")?;
|
|
|
|
// Simulate fine-tuning by modifying the pruned tensor
|
|
let updated_values: Vec<f32> = result
|
|
.pruned_tensor
|
|
.to_vec::<f32>()?
|
|
.iter()
|
|
.map(|&x| if x != 0.0 { x * 1.1 } else { x })
|
|
.collect();
|
|
|
|
let updated_tensor = Tensor::from_slice(
|
|
&updated_values,
|
|
tensor.shape().dims(),
|
|
&Device::try_default()?,
|
|
)?;
|
|
|
|
// Apply mask again to ensure sparsity is maintained
|
|
let remasked = updated_tensor.mul(&result.mask)?;
|
|
let final_sparsity = calculate_sparsity(&remasked)?;
|
|
|
|
assert!(
|
|
(final_sparsity - 0.8).abs() < 0.05,
|
|
"Sparsity should be maintained after fine-tuning: {:.2}",
|
|
final_sparsity
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_threshold_based_pruning() -> Result<()> {
|
|
let config = UnstructuredPruningConfig::new(
|
|
PruningCriterion::Threshold { threshold: 0.1 },
|
|
0.0, // Sparsity determined by threshold
|
|
);
|
|
|
|
let pruner = UnstructuredPruner::new(config)?;
|
|
|
|
// Create tensor with known distribution
|
|
let device = Device::try_default()?;
|
|
let values: Vec<f32> = (0..1000)
|
|
.map(|i| {
|
|
if i < 200 {
|
|
0.05
|
|
}
|
|
// Small values (should be pruned)
|
|
else if i < 800 {
|
|
0.15
|
|
}
|
|
// Medium values (should be kept)
|
|
else {
|
|
0.25
|
|
} // Large values (should be kept)
|
|
})
|
|
.collect();
|
|
|
|
let tensor = Tensor::from_slice(&values, &[25, 40], &device)?;
|
|
let pruned = pruner.prune_tensor(&tensor, "test")?;
|
|
|
|
let pruned_vals = pruned.to_vec::<f32>()?;
|
|
let kept_count = pruned_vals.iter().filter(|&&x| x.abs() > 1e-8).count();
|
|
|
|
// Should keep ~800 values (those >= 0.1)
|
|
assert!(
|
|
kept_count >= 750 && kept_count <= 850,
|
|
"Threshold pruning should keep ~800 values, got {}",
|
|
kept_count
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_multi_step_gradual_pruning() -> Result<()> {
|
|
let schedule = GradualPruningSchedule::exponential(
|
|
0.2, // Start at 20% sparsity
|
|
0.95, // End at 95% sparsity
|
|
500, // Over 500 steps
|
|
);
|
|
|
|
let config = UnstructuredPruningConfig::new_with_schedule(
|
|
PruningCriterion::GlobalMagnitude,
|
|
schedule,
|
|
);
|
|
|
|
let mut pruner = UnstructuredPruner::new(config)?;
|
|
let mut current_tensor = create_test_tensor(&[64, 64], 42)?;
|
|
|
|
let mut sparsity_progression = Vec::new();
|
|
let checkpoints = [0, 50, 125, 250, 375, 499];
|
|
|
|
for &step in &checkpoints {
|
|
pruner.set_current_step(step);
|
|
current_tensor = pruner.prune_tensor(¤t_tensor, "test")?;
|
|
let sparsity = calculate_sparsity(¤t_tensor)?;
|
|
sparsity_progression.push((step, sparsity));
|
|
}
|
|
|
|
// Verify progression
|
|
for i in 1..sparsity_progression.len() {
|
|
let (prev_step, prev_sparsity) = sparsity_progression[i - 1];
|
|
let (curr_step, curr_sparsity) = sparsity_progression[i];
|
|
|
|
assert!(
|
|
curr_sparsity >= prev_sparsity,
|
|
"Sparsity should increase: step {} = {:.3}, step {} = {:.3}",
|
|
prev_step,
|
|
prev_sparsity,
|
|
curr_step,
|
|
curr_sparsity
|
|
);
|
|
}
|
|
|
|
// Final sparsity should be close to target
|
|
let (_, final_sparsity) = sparsity_progression.last().unwrap();
|
|
assert!(
|
|
(final_sparsity - 0.95).abs() < 0.03,
|
|
"Final sparsity should be ~95%, got {:.2}%",
|
|
final_sparsity * 100.0
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_layerwise_sensitivity_adjustment() -> Result<()> {
|
|
// Create layer sensitivities (higher values = more sensitive = less pruning)
|
|
let mut sensitivities = HashMap::new();
|
|
sensitivities.insert("conv1.weight".to_string(), 0.9); // High sensitivity
|
|
sensitivities.insert("conv2.weight".to_string(), 0.5); // Medium sensitivity
|
|
sensitivities.insert("fc1.weight".to_string(), 0.1); // Low sensitivity
|
|
|
|
let config = UnstructuredPruningConfig::new_with_sensitivity(
|
|
PruningCriterion::LocalMagnitude,
|
|
0.7, // Base 70% sparsity
|
|
sensitivities,
|
|
);
|
|
|
|
let pruner = UnstructuredPruner::new(config)?;
|
|
let model = create_test_model(42)?;
|
|
|
|
let pruned_model = pruner.prune_model(&model)?;
|
|
|
|
// High sensitivity layer should have lower sparsity
|
|
let conv1_sparsity = calculate_sparsity(&pruned_model["conv1.weight"])?;
|
|
let conv2_sparsity = calculate_sparsity(&pruned_model["conv2.weight"])?;
|
|
let fc1_sparsity = calculate_sparsity(&pruned_model["fc1.weight"])?;
|
|
|
|
assert!(
|
|
conv1_sparsity < conv2_sparsity,
|
|
"High sensitivity layer should have lower sparsity: {:.2} vs {:.2}",
|
|
conv1_sparsity,
|
|
conv2_sparsity
|
|
);
|
|
|
|
assert!(
|
|
conv2_sparsity < fc1_sparsity,
|
|
"Medium sensitivity should have lower sparsity than low sensitivity: {:.2} vs {:.2}",
|
|
conv2_sparsity,
|
|
fc1_sparsity
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_sparse_to_sparse_continuation() -> Result<()> {
|
|
// Start with already sparse tensor
|
|
let mut initial_values = vec![0.0f32; 1024];
|
|
for i in (0..1024).step_by(3) {
|
|
// Every 3rd element is non-zero
|
|
initial_values[i] = ((i % 100) as f32) / 100.0;
|
|
}
|
|
|
|
let device = Device::try_default()?;
|
|
let sparse_tensor = Tensor::from_slice(&initial_values, &[32, 32], &device)?;
|
|
let initial_sparsity = calculate_sparsity(&sparse_tensor)?;
|
|
|
|
assert!(
|
|
initial_sparsity > 0.6,
|
|
"Initial tensor should be sparse: {:.2}",
|
|
initial_sparsity
|
|
);
|
|
|
|
// Further prune the already sparse tensor
|
|
let config = UnstructuredPruningConfig::new(
|
|
PruningCriterion::GlobalMagnitude,
|
|
0.9, // Target 90% sparsity
|
|
);
|
|
|
|
let pruner = UnstructuredPruner::new(config)?;
|
|
let further_pruned = pruner.prune_tensor(&sparse_tensor, "test")?;
|
|
let final_sparsity = calculate_sparsity(&further_pruned)?;
|
|
|
|
assert!(
|
|
final_sparsity > initial_sparsity,
|
|
"Further pruning should increase sparsity: {:.2} -> {:.2}",
|
|
initial_sparsity,
|
|
final_sparsity
|
|
);
|
|
|
|
assert!(
|
|
(final_sparsity - 0.9).abs() < 0.05,
|
|
"Should reach target sparsity: {:.2}",
|
|
final_sparsity
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_pruning_statistics_collection() -> Result<()> {
|
|
let config = UnstructuredPruningConfig::new_with_stats_collection(
|
|
PruningCriterion::GlobalMagnitude,
|
|
0.8,
|
|
true, // Collect detailed statistics
|
|
);
|
|
|
|
let pruner = UnstructuredPruner::new(config)?;
|
|
let model = create_test_model(42)?;
|
|
|
|
let _pruned_model = pruner.prune_model(&model)?;
|
|
let stats = pruner.get_statistics();
|
|
|
|
assert!(stats.total_parameters > 0);
|
|
assert!(stats.pruned_parameters > 0);
|
|
assert!(stats.compression_ratio > 1.0);
|
|
assert!(!stats.layer_statistics.is_empty());
|
|
assert!(stats.memory_saved_mb > 0.0);
|
|
|
|
// Check that statistics make sense
|
|
let expected_remaining = stats.total_parameters - stats.pruned_parameters;
|
|
assert_eq!(expected_remaining, stats.remaining_parameters);
|
|
|
|
let expected_sparsity = stats.pruned_parameters as f32 / stats.total_parameters as f32;
|
|
assert!((stats.overall_sparsity - expected_sparsity).abs() < 0.01);
|
|
|
|
Ok(())
|
|
}
|
|
}
|