482 lines
16 KiB
Rust
482 lines
16 KiB
Rust
#!/usr/bin/env rust-script
|
|
|
|
//! Standalone test for magnitude pruning implementation
|
|
//! This allows us to test the functionality without dealing with broader compilation issues
|
|
|
|
use std::collections::HashMap;
|
|
|
|
/// Result type for this standalone test
|
|
type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
|
|
|
|
/// Simple error type for testing
|
|
#[derive(Debug)]
|
|
struct TestError(String);
|
|
|
|
impl std::fmt::Display for TestError {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
write!(f, "Test error: {}", self.0)
|
|
}
|
|
}
|
|
|
|
impl std::error::Error for TestError {}
|
|
|
|
// Include our implementation directly for standalone testing
|
|
// (This is normally not recommended, but useful for TDD when dependencies have issues)
|
|
|
|
/// Configuration for magnitude-based pruning
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct PruningConfig {
|
|
/// Pruning strategy to apply
|
|
pub strategy: PruningStrategy,
|
|
/// Target sparsity ratio (0.0 to 1.0)
|
|
pub sparsity_ratio: f64,
|
|
/// Pruning granularity
|
|
pub granularity: PruningGranularity,
|
|
/// Optional pruning schedule for gradual pruning
|
|
pub schedule: Option<PruningSchedule>,
|
|
/// Random seed for deterministic pruning
|
|
pub seed: Option<u64>,
|
|
}
|
|
|
|
/// Strategy for selecting weights to prune
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub enum PruningStrategy {
|
|
/// Global magnitude pruning across all parameters
|
|
Global,
|
|
/// Layer-wise magnitude pruning within each layer
|
|
LayerWise,
|
|
/// Structured pruning (channels, filters)
|
|
Structured {
|
|
/// Type of structured pattern
|
|
pattern_type: StructuredPattern,
|
|
},
|
|
}
|
|
|
|
/// Granularity of pruning operations
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub enum PruningGranularity {
|
|
/// Element-wise pruning (unstructured)
|
|
Element,
|
|
/// Channel-wise pruning
|
|
Channel,
|
|
/// Filter-wise pruning
|
|
Filter,
|
|
/// Block-wise pruning
|
|
Block { height: usize, width: usize },
|
|
}
|
|
|
|
/// Structured pruning patterns
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub enum StructuredPattern {
|
|
/// Channel pruning
|
|
Channel,
|
|
/// Filter pruning
|
|
Filter,
|
|
/// N:M sparsity pattern
|
|
NM { n: usize, m: usize },
|
|
}
|
|
|
|
/// Pruning schedule for gradual sparsification during training
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub enum PruningSchedule {
|
|
/// Linear schedule from initial to final sparsity
|
|
Linear {
|
|
/// Initial sparsity ratio
|
|
initial_sparsity: f64,
|
|
/// Final sparsity ratio
|
|
final_sparsity: f64,
|
|
/// Duration in training steps
|
|
duration_steps: usize,
|
|
},
|
|
}
|
|
|
|
/// Binary mask for indicating which weights to keep/prune
|
|
#[derive(Debug, Clone)]
|
|
pub struct PruningMask {
|
|
/// Binary mask values (true = keep, false = prune)
|
|
mask: Vec<bool>,
|
|
/// Shape of the mask
|
|
shape: Vec<usize>,
|
|
/// Number of parameters kept
|
|
kept_parameters: usize,
|
|
/// Total number of parameters
|
|
total_parameters: usize,
|
|
}
|
|
|
|
/// Statistics about pruning results
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct PruningStats {
|
|
/// Original parameter count
|
|
pub original_params: usize,
|
|
/// Parameters after pruning
|
|
pub pruned_params: usize,
|
|
/// Achieved sparsity ratio
|
|
pub sparsity_ratio: f64,
|
|
/// Compression ratio
|
|
pub compression_ratio: f64,
|
|
/// Memory savings (in bytes, estimated)
|
|
pub memory_savings: usize,
|
|
}
|
|
|
|
/// Main struct for magnitude-based pruning operations
|
|
#[derive(Debug)]
|
|
pub struct MagnitudePruning;
|
|
|
|
/// Simple tensor for testing
|
|
#[derive(Debug, Clone)]
|
|
pub struct DenseTensor {
|
|
pub data: Vec<f32>,
|
|
pub shape: Vec<usize>,
|
|
}
|
|
|
|
impl DenseTensor {
|
|
pub fn new(data: Vec<f32>, shape: Vec<usize>) -> Self {
|
|
assert_eq!(data.len(), shape.iter().product::<usize>());
|
|
Self { data, shape }
|
|
}
|
|
|
|
pub fn shape(&self) -> &[usize] {
|
|
&self.shape
|
|
}
|
|
|
|
pub fn data(&self) -> &[f32] {
|
|
&self.data
|
|
}
|
|
}
|
|
|
|
impl PruningConfig {
|
|
/// Create global magnitude pruning configuration
|
|
pub fn global(sparsity_ratio: f64) -> Self {
|
|
Self {
|
|
strategy: PruningStrategy::Global,
|
|
sparsity_ratio,
|
|
granularity: PruningGranularity::Element,
|
|
schedule: None,
|
|
seed: None,
|
|
}
|
|
}
|
|
|
|
/// Create layer-wise magnitude pruning configuration
|
|
pub fn layerwise(sparsity_ratio: f64) -> Self {
|
|
Self {
|
|
strategy: PruningStrategy::LayerWise,
|
|
sparsity_ratio,
|
|
granularity: PruningGranularity::Element,
|
|
schedule: None,
|
|
seed: None,
|
|
}
|
|
}
|
|
|
|
/// Validate configuration parameters
|
|
pub fn validate(&self) -> Result<()> {
|
|
if !(0.0..=1.0).contains(&self.sparsity_ratio) {
|
|
return Err(Box::new(TestError("Sparsity ratio must be between 0.0 and 1.0".into())));
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl PruningMask {
|
|
/// Create new pruning mask
|
|
pub fn new(mask: Vec<bool>, shape: Vec<usize>) -> Self {
|
|
let total_parameters = shape.iter().product::<usize>();
|
|
let kept_parameters = mask.iter().filter(|&&x| x).count();
|
|
|
|
Self {
|
|
mask,
|
|
shape,
|
|
kept_parameters,
|
|
total_parameters,
|
|
}
|
|
}
|
|
|
|
/// Get mask values
|
|
pub fn mask(&self) -> &[bool] {
|
|
&self.mask
|
|
}
|
|
|
|
/// Get mask shape
|
|
pub fn shape(&self) -> &[usize] {
|
|
&self.shape
|
|
}
|
|
|
|
/// Get sparsity ratio
|
|
pub fn sparsity_ratio(&self) -> f64 {
|
|
1.0 - (self.kept_parameters as f64 / self.total_parameters as f64)
|
|
}
|
|
|
|
/// Apply mask to tensor (element-wise multiplication)
|
|
pub fn apply(&self, tensor: &DenseTensor) -> Result<DenseTensor> {
|
|
if tensor.shape != self.shape {
|
|
return Err(Box::new(TestError("Tensor and mask shapes must match".into())));
|
|
}
|
|
|
|
let masked_data: Vec<f32> = tensor.data.iter()
|
|
.zip(self.mask.iter())
|
|
.map(|(&value, &keep)| if keep { value } else { 0.0 })
|
|
.collect();
|
|
|
|
Ok(DenseTensor::new(masked_data, tensor.shape.clone()))
|
|
}
|
|
}
|
|
|
|
impl PruningStats {
|
|
/// Create new pruning statistics
|
|
pub fn new(original_params: usize, pruned_params: usize) -> Self {
|
|
let sparsity_ratio = 1.0 - (pruned_params as f64 / original_params as f64);
|
|
let compression_ratio = if pruned_params == 0 {
|
|
f64::INFINITY
|
|
} else {
|
|
original_params as f64 / pruned_params as f64
|
|
};
|
|
|
|
// Estimate memory savings (assuming 4 bytes per f32)
|
|
let memory_savings = (original_params - pruned_params) * 4;
|
|
|
|
Self {
|
|
original_params,
|
|
pruned_params,
|
|
sparsity_ratio,
|
|
compression_ratio,
|
|
memory_savings,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl MagnitudePruning {
|
|
/// Apply magnitude-based pruning to tensor
|
|
pub fn prune(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
|
|
config.validate()?;
|
|
|
|
match config.strategy {
|
|
PruningStrategy::Global => {
|
|
Self::prune_global(tensor, config)
|
|
}
|
|
PruningStrategy::LayerWise => {
|
|
Self::prune_layerwise(tensor, config)
|
|
}
|
|
PruningStrategy::Structured { .. } => {
|
|
Err(Box::new(TestError("Structured pruning not implemented yet".into())))
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Create pruning mask without modifying the original tensor
|
|
pub fn create_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
|
|
config.validate()?;
|
|
|
|
match config.strategy {
|
|
PruningStrategy::Global => {
|
|
Self::create_global_mask(tensor, config)
|
|
}
|
|
PruningStrategy::LayerWise => {
|
|
Self::create_layerwise_mask(tensor, config)
|
|
}
|
|
PruningStrategy::Structured { .. } => {
|
|
Err(Box::new(TestError("Structured mask creation not implemented yet".into())))
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Compute importance scores based on magnitude
|
|
pub fn compute_importance_scores(tensor: &DenseTensor) -> Result<Vec<f32>> {
|
|
let scores: Vec<f32> = tensor.data.iter().map(|&x| x.abs()).collect();
|
|
Ok(scores)
|
|
}
|
|
|
|
/// Analyze pruning results
|
|
pub fn analyze_pruning(original: &DenseTensor, pruned: &DenseTensor) -> Result<PruningStats> {
|
|
if original.shape != pruned.shape {
|
|
return Err(Box::new(TestError("Tensor shapes must match for analysis".into())));
|
|
}
|
|
|
|
let original_nonzeros = original.data.iter().filter(|&&x| x != 0.0).count();
|
|
let pruned_nonzeros = pruned.data.iter().filter(|&&x| x != 0.0).count();
|
|
|
|
Ok(PruningStats::new(original_nonzeros, pruned_nonzeros))
|
|
}
|
|
|
|
fn prune_global(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
|
|
let mask = Self::create_global_mask(tensor, config)?;
|
|
mask.apply(tensor)
|
|
}
|
|
|
|
fn prune_layerwise(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
|
|
let mask = Self::create_layerwise_mask(tensor, config)?;
|
|
mask.apply(tensor)
|
|
}
|
|
|
|
fn create_global_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
|
|
let total_elements = tensor.data.len();
|
|
let elements_to_keep = ((1.0 - config.sparsity_ratio) * total_elements as f64).round() as usize;
|
|
|
|
// Compute importance scores (absolute magnitudes)
|
|
let importance_scores = Self::compute_importance_scores(tensor)?;
|
|
|
|
// Create (index, score) pairs for sorting
|
|
let mut indexed_scores: Vec<(usize, f32)> = importance_scores.iter()
|
|
.enumerate()
|
|
.map(|(i, &score)| (i, score))
|
|
.collect();
|
|
|
|
// Sort by importance score in descending order (keep highest magnitude values)
|
|
indexed_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
|
|
|
|
// Create mask: true for kept elements, false for pruned elements
|
|
let mut mask_data = vec![false; total_elements];
|
|
for i in 0..elements_to_keep {
|
|
if i < indexed_scores.len() {
|
|
let (index, _) = indexed_scores[i];
|
|
mask_data[index] = true;
|
|
}
|
|
}
|
|
|
|
Ok(PruningMask::new(mask_data, tensor.shape.clone()))
|
|
}
|
|
|
|
fn create_layerwise_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
|
|
// For layer-wise pruning, we currently support 2D tensors (treating each row as a layer)
|
|
if tensor.shape.len() != 2 {
|
|
return Err(Box::new(TestError("Layer-wise pruning currently only supports 2D tensors".into())));
|
|
}
|
|
|
|
let rows = tensor.shape[0];
|
|
let cols = tensor.shape[1];
|
|
let elements_per_row = cols;
|
|
let elements_to_keep_per_row = ((1.0 - config.sparsity_ratio) * elements_per_row as f64).round() as usize;
|
|
|
|
let mut mask_data = vec![false; tensor.data.len()];
|
|
|
|
// Process each row independently
|
|
for row in 0..rows {
|
|
let row_start = row * cols;
|
|
let row_end = row_start + cols;
|
|
|
|
// Get magnitudes for this row
|
|
let row_scores: Vec<(usize, f32)> = (row_start..row_end)
|
|
.map(|i| (i, tensor.data[i].abs()))
|
|
.collect();
|
|
|
|
// Sort by magnitude in descending order
|
|
let mut sorted_scores = row_scores;
|
|
sorted_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
|
|
|
|
// Keep the top elements_to_keep_per_row elements in this row
|
|
for i in 0..elements_to_keep_per_row.min(sorted_scores.len()) {
|
|
let (original_index, _) = sorted_scores[i];
|
|
mask_data[original_index] = true;
|
|
}
|
|
}
|
|
|
|
Ok(PruningMask::new(mask_data, tensor.shape.clone()))
|
|
}
|
|
}
|
|
|
|
/// Test utilities
|
|
mod test_utils {
|
|
use super::*;
|
|
|
|
/// Create a simple 2D test tensor with known values
|
|
pub fn create_test_tensor_2d() -> DenseTensor {
|
|
// 3x4 tensor with specific values for predictable pruning results
|
|
let data = vec![
|
|
0.1, 0.9, 0.2, 0.8, // Row 0: magnitudes [0.1, 0.9, 0.2, 0.8]
|
|
0.3, 0.7, 0.4, 0.6, // Row 1: magnitudes [0.3, 0.7, 0.4, 0.6]
|
|
0.05, 0.95, 0.15, 0.85 // Row 2: magnitudes [0.05, 0.95, 0.15, 0.85]
|
|
];
|
|
DenseTensor::new(data, vec![3, 4])
|
|
}
|
|
}
|
|
|
|
fn main() -> Result<()> {
|
|
println!("Running magnitude pruning standalone tests...");
|
|
|
|
// Test 1: Basic configuration
|
|
println!("\n=== Test 1: Basic Configuration ===");
|
|
let config = PruningConfig::global(0.5);
|
|
println!("Global config created: {:?}", config);
|
|
config.validate()?;
|
|
println!("✓ Configuration validation passed");
|
|
|
|
// Test 2: Importance score computation
|
|
println!("\n=== Test 2: Importance Score Computation ===");
|
|
let tensor = test_utils::create_test_tensor_2d();
|
|
println!("Test tensor: {:?}", tensor);
|
|
let scores = MagnitudePruning::compute_importance_scores(&tensor)?;
|
|
println!("Importance scores: {:?}", scores);
|
|
|
|
let expected_scores = vec![0.1, 0.9, 0.2, 0.8, 0.3, 0.7, 0.4, 0.6, 0.05, 0.95, 0.15, 0.85];
|
|
assert_eq!(scores, expected_scores);
|
|
println!("✓ Importance scores computation passed");
|
|
|
|
// Test 3: Global magnitude pruning
|
|
println!("\n=== Test 3: Global Magnitude Pruning ===");
|
|
let config = PruningConfig::global(0.5); // 50% sparsity
|
|
let mask = MagnitudePruning::create_mask(&tensor, &config)?;
|
|
println!("Global mask sparsity ratio: {:.2}", mask.sparsity_ratio());
|
|
println!("Kept parameters: {}/{}", mask.kept_parameters, mask.total_parameters);
|
|
|
|
let pruned = MagnitudePruning::prune(&tensor, &config)?;
|
|
println!("Pruned tensor data: {:?}", pruned.data);
|
|
|
|
let nonzero_count = pruned.data.iter().filter(|&&x| x != 0.0).count();
|
|
println!("Non-zero count after pruning: {}", nonzero_count);
|
|
assert_eq!(nonzero_count, 6); // Should keep 6 out of 12 (50% sparsity)
|
|
|
|
// Verify the largest magnitude values are kept: [0.95, 0.9, 0.85, 0.8, 0.7, 0.6]
|
|
let mut nonzero_values: Vec<f32> = pruned.data.iter().filter(|&&x| x != 0.0).copied().collect();
|
|
nonzero_values.sort_by(|a, b| b.abs().total_cmp(&a.abs()));
|
|
let expected_kept = vec![0.95, 0.9, 0.85, 0.8, 0.7, 0.6];
|
|
assert_eq!(nonzero_values, expected_kept);
|
|
println!("✓ Global magnitude pruning passed");
|
|
|
|
// Test 4: Layer-wise magnitude pruning
|
|
println!("\n=== Test 4: Layer-wise Magnitude Pruning ===");
|
|
let config = PruningConfig::layerwise(0.5);
|
|
let pruned = MagnitudePruning::prune(&tensor, &config)?;
|
|
println!("Layer-wise pruned tensor data: {:?}", pruned.data);
|
|
|
|
// With layer-wise pruning on a 2D tensor, each row should have 50% sparsity
|
|
// Row 0: keep 2 largest [0.9, 0.8], prune [0.1, 0.2]
|
|
// Row 1: keep 2 largest [0.7, 0.6], prune [0.3, 0.4]
|
|
// Row 2: keep 2 largest [0.95, 0.85], prune [0.05, 0.15]
|
|
let expected_data = vec![
|
|
0.0, 0.9, 0.0, 0.8,
|
|
0.0, 0.7, 0.0, 0.6,
|
|
0.0, 0.95, 0.0, 0.85
|
|
];
|
|
assert_eq!(pruned.data, expected_data);
|
|
println!("✓ Layer-wise magnitude pruning passed");
|
|
|
|
// Test 5: Pruning analysis
|
|
println!("\n=== Test 5: Pruning Analysis ===");
|
|
let stats = MagnitudePruning::analyze_pruning(&tensor, &pruned)?;
|
|
println!("Pruning statistics: {:?}", stats);
|
|
assert_eq!(stats.original_params, 12);
|
|
assert_eq!(stats.pruned_params, 6);
|
|
assert!((stats.sparsity_ratio - 0.5).abs() < f64::EPSILON);
|
|
println!("✓ Pruning analysis passed");
|
|
|
|
// Test 6: Mask application
|
|
println!("\n=== Test 6: Mask Application ===");
|
|
let tensor_for_mask = test_utils::create_test_tensor_2d();
|
|
let mask_data = vec![
|
|
true, false, true, false, // Keep [0.1, _, 0.2, _]
|
|
false, true, false, true, // Keep [_, 0.7, _, 0.6]
|
|
true, true, false, false // Keep [0.05, 0.95, _, _]
|
|
];
|
|
let mask = PruningMask::new(mask_data, vec![3, 4]);
|
|
|
|
let masked_tensor = mask.apply(&tensor_for_mask)?;
|
|
let expected_masked = vec![
|
|
0.1, 0.0, 0.2, 0.0,
|
|
0.0, 0.7, 0.0, 0.6,
|
|
0.05, 0.95, 0.0, 0.0
|
|
];
|
|
assert_eq!(masked_tensor.data, expected_masked);
|
|
println!("✓ Mask application passed");
|
|
|
|
println!("\n🎉 All tests passed! Magnitude pruning implementation is working correctly.");
|
|
|
|
Ok(())
|
|
} |