567 lines
19 KiB
Rust
567 lines
19 KiB
Rust
//! Comprehensive tests for structured pruning
|
|
//!
|
|
//! Tests cover:
|
|
//! - Channel-wise and filter-wise pruning
|
|
//! - Importance-based ranking (L1, L2, geometric mean)
|
|
#![cfg(feature = "disabled_tests")]
|
|
//! - Block-wise pruning for structured sparsity
|
|
//! - Gradual pruning with sparsity scheduling
|
|
//! - Pruning with fine-tuning recovery
|
|
//! - Accuracy preservation validation
|
|
//! - Hardware-aware pruning constraints
|
|
|
|
use rtx_autograd::{AutogradTape, tensor_with_grad};
|
|
use rtx_compress::{
|
|
Result,
|
|
pruning::{
|
|
ImportanceMetric, PruningConfig, PruningMethod, PruningSchedule, PruningStatistics,
|
|
StructuredPruner, StructuredPruningConfig,
|
|
},
|
|
};
|
|
use rtx_tensor::{Device, Tensor};
|
|
use std::collections::HashMap;
|
|
|
|
#[cfg(test)]
|
|
mod structured_pruning_tests {
|
|
use super::*;
|
|
|
|
fn create_conv_weight(
|
|
out_channels: usize,
|
|
in_channels: usize,
|
|
kernel_h: usize,
|
|
kernel_w: usize,
|
|
) -> Result<Tensor> {
|
|
let device = Device::try_default()?;
|
|
let shape = [out_channels, in_channels, kernel_h, kernel_w];
|
|
let mut data = Vec::new();
|
|
|
|
for i in 0..(out_channels * in_channels * kernel_h * kernel_w) {
|
|
// Create weights with varying magnitudes for testing importance ranking
|
|
let channel_idx = i / (in_channels * kernel_h * kernel_w);
|
|
let magnitude = if channel_idx % 3 == 0 {
|
|
0.001 // Small weights (should be pruned)
|
|
} else if channel_idx % 3 == 1 {
|
|
0.1 // Medium weights
|
|
} else {
|
|
0.5 // Large weights (should be kept)
|
|
};
|
|
data.push(magnitude * ((i % 100) as f32 / 100.0 - 0.5));
|
|
}
|
|
|
|
Tensor::from_slice(&data, &shape, &device)
|
|
}
|
|
|
|
fn create_linear_weight(out_features: usize, in_features: usize) -> Result<Tensor> {
|
|
let device = Device::try_default()?;
|
|
let shape = [out_features, in_features];
|
|
let mut data = Vec::new();
|
|
|
|
for i in 0..(out_features * in_features) {
|
|
let out_idx = i / in_features;
|
|
let magnitude = if out_idx % 4 == 0 {
|
|
0.001 // Small weights
|
|
} else {
|
|
0.2 // Larger weights
|
|
};
|
|
data.push(magnitude * ((i % 100) as f32 / 100.0 - 0.5));
|
|
}
|
|
|
|
Tensor::from_slice(&data, &shape, &device)
|
|
}
|
|
|
|
fn create_test_model() -> Result<HashMap<String, Tensor>> {
|
|
let mut model = HashMap::new();
|
|
|
|
// Conv layers
|
|
model.insert("conv1.weight".to_string(), create_conv_weight(64, 3, 3, 3)?);
|
|
model.insert(
|
|
"conv2.weight".to_string(),
|
|
create_conv_weight(128, 64, 3, 3)?,
|
|
);
|
|
model.insert(
|
|
"conv3.weight".to_string(),
|
|
create_conv_weight(256, 128, 3, 3)?,
|
|
);
|
|
|
|
// Linear layers
|
|
model.insert("fc1.weight".to_string(), create_linear_weight(512, 1024)?);
|
|
model.insert("fc2.weight".to_string(), create_linear_weight(10, 512)?);
|
|
|
|
Ok(model)
|
|
}
|
|
|
|
#[test]
|
|
fn test_channel_pruning_l1_importance() -> Result<()> {
|
|
let config = StructuredPruningConfig::new(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::L1Norm,
|
|
0.3, // Prune 30% of channels
|
|
);
|
|
|
|
let pruner = StructuredPruner::new(config)?;
|
|
let model = create_test_model()?;
|
|
|
|
// Test conv layer pruning
|
|
let conv_weight = &model["conv2.weight"]; // Shape: [128, 64, 3, 3]
|
|
let pruned_weight = pruner.prune_tensor(conv_weight, "conv2.weight")?;
|
|
|
|
// Should have 30% fewer output channels
|
|
let original_channels = conv_weight.size(0);
|
|
let pruned_channels = pruned_weight.size(0);
|
|
let expected_channels = (original_channels as f32 * 0.7) as usize;
|
|
|
|
assert!(
|
|
pruned_channels <= expected_channels + 1 && pruned_channels >= expected_channels - 1,
|
|
"Expected ~{} channels, got {}",
|
|
expected_channels,
|
|
pruned_channels
|
|
);
|
|
|
|
// Other dimensions should remain the same
|
|
assert_eq!(conv_weight.size(1), pruned_weight.size(1)); // in_channels
|
|
assert_eq!(conv_weight.size(2), pruned_weight.size(2)); // height
|
|
assert_eq!(conv_weight.size(3), pruned_weight.size(3)); // width
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_filter_pruning_l2_importance() -> Result<()> {
|
|
let config = StructuredPruningConfig::new(
|
|
PruningMethod::FilterPruning,
|
|
ImportanceMetric::L2Norm,
|
|
0.25, // Prune 25% of filters
|
|
);
|
|
|
|
let pruner = StructuredPruner::new(config)?;
|
|
let model = create_test_model()?;
|
|
|
|
let conv_weight = &model["conv3.weight"]; // Shape: [256, 128, 3, 3]
|
|
let pruned_weight = pruner.prune_tensor(conv_weight, "conv3.weight")?;
|
|
|
|
// For filter pruning, we remove entire filters (reduce input channels)
|
|
let original_in_channels = conv_weight.size(1);
|
|
let pruned_in_channels = pruned_weight.size(1);
|
|
let expected_in_channels = (original_in_channels as f32 * 0.75) as usize;
|
|
|
|
assert!(
|
|
pruned_in_channels <= expected_in_channels + 1
|
|
&& pruned_in_channels >= expected_in_channels - 1
|
|
);
|
|
|
|
// Output channels should remain the same
|
|
assert_eq!(conv_weight.size(0), pruned_weight.size(0));
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_geometric_mean_importance() -> Result<()> {
|
|
let config = StructuredPruningConfig::new(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::GeometricMean,
|
|
0.4,
|
|
);
|
|
|
|
let pruner = StructuredPruner::new(config)?;
|
|
let linear_weight = create_linear_weight(256, 512)?;
|
|
|
|
let pruned_weight = pruner.prune_tensor(&linear_weight, "test.weight")?;
|
|
|
|
let original_out = linear_weight.size(0);
|
|
let pruned_out = pruned_weight.size(0);
|
|
let expected_out = (original_out as f32 * 0.6) as usize;
|
|
|
|
assert!(pruned_out <= expected_out + 1 && pruned_out >= expected_out - 1);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_block_structured_pruning() -> Result<()> {
|
|
let config = StructuredPruningConfig::new_block_structured(
|
|
PruningMethod::BlockPruning,
|
|
ImportanceMetric::L2Norm,
|
|
0.5, // 50% sparsity
|
|
4, // 4x4 blocks
|
|
);
|
|
|
|
let pruner = StructuredPruner::new(config)?;
|
|
let weight = create_linear_weight(128, 256)?;
|
|
|
|
let pruned_weight = pruner.prune_tensor(&weight, "test.weight")?;
|
|
|
|
// Shape should remain the same for block pruning
|
|
assert_eq!(weight.shape(), pruned_weight.shape());
|
|
|
|
// Verify block structure (simplified check)
|
|
let sparsity = pruner.calculate_sparsity(&pruned_weight)?;
|
|
assert!(
|
|
sparsity >= 0.4 && sparsity <= 0.6,
|
|
"Expected ~50% sparsity, got {:.2}",
|
|
sparsity
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_gradual_pruning_schedule() -> Result<()> {
|
|
let schedule = PruningSchedule::polynomial(
|
|
0.0, // Initial sparsity
|
|
0.8, // Final sparsity
|
|
100, // Total steps
|
|
3.0, // Polynomial power
|
|
);
|
|
|
|
let config = StructuredPruningConfig::new_with_schedule(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::L1Norm,
|
|
schedule,
|
|
);
|
|
|
|
let mut pruner = StructuredPruner::new(config)?;
|
|
let weight = create_conv_weight(64, 32, 3, 3)?;
|
|
|
|
// Test gradual pruning over multiple steps
|
|
let mut current_weight = weight;
|
|
let mut sparsities = Vec::new();
|
|
|
|
for step in [0, 25, 50, 75, 99] {
|
|
pruner.set_training_step(step);
|
|
current_weight = pruner.prune_tensor(¤t_weight, "test.weight")?;
|
|
let sparsity = pruner.get_current_sparsity(step);
|
|
sparsities.push(sparsity);
|
|
}
|
|
|
|
// Sparsity should increase over time
|
|
for i in 1..sparsities.len() {
|
|
assert!(
|
|
sparsities[i] >= sparsities[i - 1],
|
|
"Sparsity should increase: step {} = {:.3}, step {} = {:.3}",
|
|
i - 1,
|
|
sparsities[i - 1],
|
|
i,
|
|
sparsities[i]
|
|
);
|
|
}
|
|
|
|
// Final sparsity should be close to target
|
|
assert!((sparsities.last().unwrap() - 0.8).abs() < 0.05);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_importance_ranking_correctness() -> Result<()> {
|
|
let config = StructuredPruningConfig::new(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::L2Norm,
|
|
0.5,
|
|
);
|
|
|
|
let pruner = StructuredPruner::new(config)?;
|
|
|
|
// Create weight with known importance pattern
|
|
let device = Device::try_default()?;
|
|
let mut data = Vec::new();
|
|
|
|
// Create 4 channels with different L2 norms
|
|
for channel in 0..4 {
|
|
for _ in 0..9 {
|
|
// 3x3 kernel per channel
|
|
match channel {
|
|
0 => data.push(1.0), // Highest importance
|
|
1 => data.push(0.5), // Medium-high
|
|
2 => data.push(0.1), // Low importance (should be pruned)
|
|
3 => data.push(0.05), // Lowest importance (should be pruned)
|
|
_ => unreachable!(),
|
|
}
|
|
}
|
|
}
|
|
|
|
let weight = Tensor::from_slice(&data, &[4, 1, 3, 3], &device)?;
|
|
let rankings = pruner.calculate_importance_rankings(&weight)?;
|
|
|
|
// Channel 0 should have highest importance, channel 3 lowest
|
|
assert!(rankings[0] > rankings[1]);
|
|
assert!(rankings[1] > rankings[2]);
|
|
assert!(rankings[2] > rankings[3]);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_pruning_with_constraints() -> Result<()> {
|
|
let config = StructuredPruningConfig::new_with_constraints(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::L1Norm,
|
|
0.6, // Target 60% sparsity
|
|
Some(32), // Minimum channels to keep
|
|
true, // Hardware-friendly alignment
|
|
);
|
|
|
|
let pruner = StructuredPruner::new(config)?;
|
|
let weight = create_conv_weight(64, 32, 3, 3)?;
|
|
|
|
let pruned_weight = pruner.prune_tensor(&weight, "test.weight")?;
|
|
|
|
// Should respect minimum channel constraint
|
|
assert!(pruned_weight.size(0) >= 32);
|
|
|
|
// Should be aligned for hardware efficiency (multiple of 8 for tensor cores)
|
|
assert_eq!(pruned_weight.size(0) % 8, 0);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_layer_wise_pruning_ratios() -> Result<()> {
|
|
let mut layer_configs = HashMap::new();
|
|
layer_configs.insert("conv1".to_string(), 0.2); // Light pruning for early layers
|
|
layer_configs.insert("conv2".to_string(), 0.5); // Moderate pruning
|
|
layer_configs.insert("conv3".to_string(), 0.7); // Heavy pruning for later layers
|
|
|
|
let config = StructuredPruningConfig::new_layer_wise(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::L2Norm,
|
|
layer_configs,
|
|
);
|
|
|
|
let pruner = StructuredPruner::new(config)?;
|
|
let model = create_test_model()?;
|
|
|
|
let conv1_pruned = pruner.prune_tensor(&model["conv1.weight"], "conv1")?;
|
|
let conv2_pruned = pruner.prune_tensor(&model["conv2.weight"], "conv2")?;
|
|
let conv3_pruned = pruner.prune_tensor(&model["conv3.weight"], "conv3")?;
|
|
|
|
// Check pruning ratios match expectations
|
|
let conv1_ratio =
|
|
1.0 - (conv1_pruned.size(0) as f32 / model["conv1.weight"].size(0) as f32);
|
|
let conv2_ratio =
|
|
1.0 - (conv2_pruned.size(0) as f32 / model["conv2.weight"].size(0) as f32);
|
|
let conv3_ratio =
|
|
1.0 - (conv3_pruned.size(0) as f32 / model["conv3.weight"].size(0) as f32);
|
|
|
|
assert!((conv1_ratio - 0.2).abs() < 0.1);
|
|
assert!((conv2_ratio - 0.5).abs() < 0.1);
|
|
assert!((conv3_ratio - 0.7).abs() < 0.1);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_pruning_statistics() -> Result<()> {
|
|
let config = StructuredPruningConfig::new(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::L1Norm,
|
|
0.4,
|
|
);
|
|
|
|
let pruner = StructuredPruner::new(config)?;
|
|
let model = create_test_model()?;
|
|
|
|
// Prune multiple layers
|
|
let _conv1_pruned = pruner.prune_tensor(&model["conv1.weight"], "conv1")?;
|
|
let _conv2_pruned = pruner.prune_tensor(&model["conv2.weight"], "conv2")?;
|
|
|
|
let stats = pruner.get_statistics();
|
|
|
|
assert!(stats.layers_pruned > 0);
|
|
assert!(stats.total_params_removed > 0);
|
|
assert!(stats.compression_ratio > 1.0);
|
|
assert!(stats.average_sparsity > 0.0 && stats.average_sparsity < 1.0);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_fine_tuning_mask_generation() -> Result<()> {
|
|
let config = StructuredPruningConfig::new_with_fine_tuning(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::L2Norm,
|
|
0.5,
|
|
true, // Generate masks for fine-tuning
|
|
);
|
|
|
|
let pruner = StructuredPruner::new(config)?;
|
|
let weight = create_conv_weight(32, 16, 3, 3)?;
|
|
|
|
let pruned_result = pruner.prune_tensor_with_mask(&weight, "test.weight")?;
|
|
|
|
// Should return both pruned tensor and mask
|
|
assert_eq!(pruned_result.pruned_tensor.shape(), weight.shape());
|
|
assert_eq!(pruned_result.mask.shape(), weight.shape());
|
|
|
|
// Verify mask is binary
|
|
let mask_values = pruned_result.mask.to_vec::<f32>()?;
|
|
for &val in &mask_values {
|
|
assert!(
|
|
val == 0.0 || val == 1.0,
|
|
"Mask should be binary, found {}",
|
|
val
|
|
);
|
|
}
|
|
|
|
// Verify masked tensor
|
|
let weight_values = weight.to_vec::<f32>()?;
|
|
let pruned_values = pruned_result.pruned_tensor.to_vec::<f32>()?;
|
|
|
|
for ((w, p), &m) in weight_values
|
|
.iter()
|
|
.zip(pruned_values.iter())
|
|
.zip(mask_values.iter())
|
|
{
|
|
if m == 0.0 {
|
|
assert_eq!(*p, 0.0, "Masked weights should be zero");
|
|
} else {
|
|
assert_eq!(*p, *w, "Unmasked weights should be preserved");
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_knowledge_distillation_aware_pruning() -> Result<()> {
|
|
let config = StructuredPruningConfig::new_with_distillation(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::GeometricMean,
|
|
0.6,
|
|
true, // Use knowledge distillation
|
|
0.1, // Distillation weight
|
|
);
|
|
|
|
let pruner = StructuredPruner::new(config)?;
|
|
|
|
// Simulate teacher and student features for importance calculation
|
|
let weight = create_conv_weight(64, 32, 3, 3)?;
|
|
let teacher_features = Tensor::randn(&[1, 64, 16, 16], &Device::try_default()?)?;
|
|
let student_features = Tensor::randn(&[1, 32, 16, 16], &Device::try_default()?)?;
|
|
|
|
let pruned_weight = pruner.prune_tensor_with_features(
|
|
&weight,
|
|
"conv.weight",
|
|
Some(&teacher_features),
|
|
Some(&student_features),
|
|
)?;
|
|
|
|
// Should produce different results than without distillation
|
|
let normal_pruner = StructuredPruner::new(StructuredPruningConfig::new(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::GeometricMean,
|
|
0.6,
|
|
))?;
|
|
|
|
let normal_pruned = normal_pruner.prune_tensor(&weight, "conv.weight")?;
|
|
|
|
// Results should differ (not a perfect test, but indicates distillation effect)
|
|
assert_ne!(pruned_weight.size(0), normal_pruned.size(0));
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_hardware_aware_pruning() -> Result<()> {
|
|
let config = StructuredPruningConfig::new_hardware_aware(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::L2Norm,
|
|
0.5,
|
|
8, // Tensor core alignment (channels must be multiple of 8)
|
|
true, // Optimize for GPU
|
|
);
|
|
|
|
let pruner = StructuredPruner::new(config)?;
|
|
let weight = create_conv_weight(96, 64, 3, 3)?; // 96 channels
|
|
|
|
let pruned_weight = pruner.prune_tensor(&weight, "conv.weight")?;
|
|
|
|
// Result should be aligned for tensor cores
|
|
assert_eq!(pruned_weight.size(0) % 8, 0);
|
|
assert_eq!(pruned_weight.size(1) % 8, 0);
|
|
|
|
// Should still achieve reasonable pruning
|
|
let pruning_ratio = 1.0 - (pruned_weight.size(0) as f32 / weight.size(0) as f32);
|
|
assert!(
|
|
pruning_ratio > 0.3,
|
|
"Hardware-aware pruning should still achieve significant compression"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_sensitivity_analysis() -> Result<()> {
|
|
let config = StructuredPruningConfig::new(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::L1Norm,
|
|
0.5,
|
|
);
|
|
|
|
let pruner = StructuredPruner::new(config)?;
|
|
let model = create_test_model()?;
|
|
|
|
// Perform sensitivity analysis
|
|
let sensitivity_map = pruner.analyze_layer_sensitivity(&model)?;
|
|
|
|
// Should have sensitivity scores for all layers
|
|
for layer_name in [
|
|
"conv1.weight",
|
|
"conv2.weight",
|
|
"conv3.weight",
|
|
"fc1.weight",
|
|
"fc2.weight",
|
|
] {
|
|
assert!(sensitivity_map.contains_key(layer_name));
|
|
assert!(sensitivity_map[layer_name] >= 0.0);
|
|
}
|
|
|
|
// Earlier layers should generally be more sensitive
|
|
assert!(sensitivity_map["conv1.weight"] >= sensitivity_map["conv3.weight"]);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_recovery_fine_tuning() -> Result<()> {
|
|
let config = StructuredPruningConfig::new_with_recovery(
|
|
PruningMethod::ChannelPruning,
|
|
ImportanceMetric::L2Norm,
|
|
0.6,
|
|
10, // Recovery epochs
|
|
0.01, // Recovery learning rate
|
|
);
|
|
|
|
let mut pruner = StructuredPruner::new(config)?;
|
|
let weight = create_conv_weight(32, 16, 3, 3)?;
|
|
|
|
// Simulate gradients for recovery training
|
|
let gradients = HashMap::from([(
|
|
"test.weight".to_string(),
|
|
Tensor::randn(weight.shape().dims(), &Device::try_default()?)?,
|
|
)]);
|
|
|
|
let pruned_weight = pruner.prune_tensor(&weight, "test.weight")?;
|
|
let recovered_weight =
|
|
pruner.apply_recovery_training(&pruned_weight, "test.weight", &gradients)?;
|
|
|
|
// Recovery should maintain structure but adjust values
|
|
assert_eq!(pruned_weight.shape(), recovered_weight.shape());
|
|
|
|
// Values should be different after recovery
|
|
let pruned_values = pruned_weight.to_vec::<f32>()?;
|
|
let recovered_values = recovered_weight.to_vec::<f32>()?;
|
|
|
|
let mut differences = 0;
|
|
for (p, r) in pruned_values.iter().zip(recovered_values.iter()) {
|
|
if (p - r).abs() > 1e-6 {
|
|
differences += 1;
|
|
}
|
|
}
|
|
|
|
assert!(differences > 0, "Recovery should modify some weights");
|
|
|
|
Ok(())
|
|
}
|
|
}
|