torsh_ffi/
model_optimization.rs

1//! Model Optimization Techniques for Compression and Acceleration
2//!
3//! This module provides advanced model optimization techniques including pruning,
4//! knowledge distillation, neural architecture search, and operator fusion for
5//! deploying efficient deep learning models on edge devices.
6//!
7//! # Features
8//!
9//! - **Structured Pruning**: Remove entire channels/filters
10//! - **Unstructured Pruning**: Remove individual weights
11//! - **Magnitude-based Pruning**: Prune weights below threshold
12//! - **Gradient-based Pruning**: Prune based on gradient magnitude
13//! - **Knowledge Distillation**: Transfer knowledge from large teacher to small student
14//! - **Operator Fusion**: Combine operations for efficiency
15//! - **Layer Fusion**: Merge consecutive layers
16//!
17//! # Architecture
18//!
19//! ```text
20//! ┌─────────────────────────────────────────────────┐
21//! │          Original Dense Model                   │
22//! │     (100% parameters, baseline speed)           │
23//! └───────────────────┬─────────────────────────────┘
24//!                     │
25//!         ┌───────────▼──────────┐
26//!         │  Optimization        │
27//!         │                      │
28//!    ┌────┴────┐          ┌─────┴────┐
29//!    │ Pruning │          │Distill   │
30//!    └────┬────┘          └─────┬────┘
31//!         │                     │
32//!         └──────────┬──────────┘
33//!                    │
34//!     ┌──────────────▼──────────────┐
35//!     │  Sparse Optimized Model     │
36//!     │  (20-50% params, 2-5x speed)│
37//!     └──────────────┬──────────────┘
38//!                    │
39//!     ┌──────────────▼──────────────┐
40//!     │  Operator Fusion            │
41//!     │  • Conv + BN + ReLU → Fused│
42//!     │  • Linear + Bias → Fused   │
43//!     └──────────────┬──────────────┘
44//!                    │
45//!     ┌──────────────▼──────────────┐
46//!     │  Production-Ready Model     │
47//!     │  (Sparse + Fused + Fast)    │
48//!     └─────────────────────────────┘
49//! ```
50//!
51//! # Quick Start
52//!
53//! ## Magnitude-Based Pruning
54//!
55//! ```rust,ignore
56//! use torsh_ffi::model_optimization::{PruningConfig, Pruner, PruningStrategy};
57//!
58//! // Configure pruning
59//! let config = PruningConfig::new()
60//!     .with_strategy(PruningStrategy::Magnitude)
61//!     .with_sparsity(0.5)  // Remove 50% of weights
62//!     .with_schedule(PruningSchedule::Gradual { steps: 1000 });
63//!
64//! // Prune model
65//! let pruner = Pruner::new(config);
66//! let pruned_model = pruner.prune(&model)?;
67//!
68//! println!("Sparsity: {:.1}%", pruned_model.sparsity() * 100.0);
69//! ```
70//!
71//! ## Knowledge Distillation
72//!
73//! ```rust,ignore
74//! use torsh_ffi::model_optimization::{DistillationConfig, Distiller};
75//!
76//! // Configure distillation
77//! let config = DistillationConfig::new()
78//!     .with_temperature(3.0)
79//!     .with_alpha(0.7);  // 70% soft targets, 30% hard targets
80//!
81//! // Train student with teacher
82//! let distiller = Distiller::new(teacher_model, student_model, config);
83//!
84//! for epoch in 0..epochs {
85//!     let loss = distiller.train_step(&input, &target)?;
86//!     println!("Epoch {}: Loss {:.4}", epoch, loss);
87//! }
88//! ```
89
90use serde::{Deserialize, Serialize};
91
92/// Pruning strategy
93#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
94pub enum PruningStrategy {
95    /// Magnitude-based pruning (remove smallest weights)
96    Magnitude,
97    /// Gradient-based pruning (remove weights with small gradients)
98    Gradient,
99    /// Random pruning (baseline)
100    Random,
101    /// Structured pruning (remove entire channels/filters)
102    Structured,
103    /// L1-norm based pruning
104    L1Norm,
105    /// L2-norm based pruning
106    L2Norm,
107}
108
109/// Pruning schedule
110#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
111pub enum PruningSchedule {
112    /// One-shot pruning (prune all at once)
113    OneShot,
114    /// Gradual pruning over training
115    Gradual {
116        /// Number of steps to reach target sparsity
117        steps: usize,
118    },
119    /// Iterative pruning (prune, train, prune, train...)
120    Iterative {
121        /// Number of iterations
122        iterations: usize,
123        /// Training steps per iteration
124        train_steps: usize,
125    },
126}
127
128/// Pruning configuration
129#[derive(Debug, Clone, Serialize, Deserialize)]
130pub struct PruningConfig {
131    /// Pruning strategy
132    pub strategy: PruningStrategy,
133    /// Target sparsity (0.0 to 1.0)
134    pub sparsity: f32,
135    /// Pruning schedule
136    pub schedule: PruningSchedule,
137    /// Layers to prune (empty = all layers)
138    pub layers_to_prune: Vec<String>,
139    /// Layers to skip
140    pub layers_to_skip: Vec<String>,
141    /// Whether to use magnitude pruning
142    pub use_magnitude: bool,
143    /// Whether to update masks during training
144    pub update_masks: bool,
145}
146
147impl PruningConfig {
148    /// Create a new pruning configuration
149    pub fn new() -> Self {
150        Self {
151            strategy: PruningStrategy::Magnitude,
152            sparsity: 0.5,
153            schedule: PruningSchedule::OneShot,
154            layers_to_prune: Vec::new(),
155            layers_to_skip: vec!["output".to_string()], // Don't prune output layer
156            use_magnitude: true,
157            update_masks: false,
158        }
159    }
160
161    /// Set pruning strategy
162    pub fn with_strategy(mut self, strategy: PruningStrategy) -> Self {
163        self.strategy = strategy;
164        self
165    }
166
167    /// Set target sparsity
168    pub fn with_sparsity(mut self, sparsity: f32) -> Self {
169        self.sparsity = sparsity.clamp(0.0, 1.0);
170        self
171    }
172
173    /// Set pruning schedule
174    pub fn with_schedule(mut self, schedule: PruningSchedule) -> Self {
175        self.schedule = schedule;
176        self
177    }
178
179    /// Add layer to prune
180    pub fn prune_layer(mut self, layer: String) -> Self {
181        self.layers_to_prune.push(layer);
182        self
183    }
184
185    /// Add layer to skip
186    pub fn skip_layer(mut self, layer: String) -> Self {
187        self.layers_to_skip.push(layer);
188        self
189    }
190}
191
192impl Default for PruningConfig {
193    fn default() -> Self {
194        Self::new()
195    }
196}
197
198/// Pruning mask for a layer
199#[derive(Debug, Clone, Serialize, Deserialize)]
200pub struct PruningMask {
201    /// Binary mask (1 = keep, 0 = prune)
202    pub mask: Vec<bool>,
203    /// Current sparsity of this mask
204    pub sparsity: f32,
205    /// Layer name
206    pub layer_name: String,
207}
208
209impl PruningMask {
210    /// Create a new pruning mask
211    pub fn new(size: usize, layer_name: String) -> Self {
212        Self {
213            mask: vec![true; size], // Start with all weights kept
214            sparsity: 0.0,
215            layer_name,
216        }
217    }
218
219    /// Apply magnitude-based pruning
220    ///
221    /// # Arguments
222    /// * `weights` - Weight values
223    /// * `target_sparsity` - Target sparsity (0.0 to 1.0)
224    pub fn apply_magnitude_pruning(&mut self, weights: &[f32], target_sparsity: f32) {
225        // Get absolute magnitudes
226        let mut magnitudes: Vec<(usize, f32)> = weights
227            .iter()
228            .enumerate()
229            .map(|(i, &w)| (i, w.abs()))
230            .collect();
231
232        // Sort by magnitude (ascending)
233        magnitudes.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
234
235        // Calculate number of weights to prune
236        let num_to_prune = (weights.len() as f32 * target_sparsity) as usize;
237
238        // Prune smallest magnitude weights
239        for i in 0..num_to_prune.min(magnitudes.len()) {
240            let idx = magnitudes[i].0;
241            self.mask[idx] = false;
242        }
243
244        // Update sparsity
245        self.update_sparsity();
246    }
247
248    /// Apply random pruning (for baseline comparison)
249    pub fn apply_random_pruning(&mut self, target_sparsity: f32) {
250        let num_to_prune = (self.mask.len() as f32 * target_sparsity) as usize;
251
252        for _ in 0..num_to_prune {
253            let idx = fastrand::usize(..self.mask.len());
254            self.mask[idx] = false;
255        }
256
257        self.update_sparsity();
258    }
259
260    /// Apply the mask to weights (zero out pruned weights)
261    pub fn apply_to_weights(&self, weights: &mut [f32]) {
262        for (i, &keep) in self.mask.iter().enumerate() {
263            if !keep && i < weights.len() {
264                weights[i] = 0.0;
265            }
266        }
267    }
268
269    /// Update sparsity calculation
270    pub fn update_sparsity(&mut self) {
271        let pruned = self.mask.iter().filter(|&&x| !x).count();
272        self.sparsity = pruned as f32 / self.mask.len() as f32;
273    }
274
275    /// Get number of pruned parameters
276    pub fn num_pruned(&self) -> usize {
277        self.mask.iter().filter(|&&x| !x).count()
278    }
279
280    /// Get number of kept parameters
281    pub fn num_kept(&self) -> usize {
282        self.mask.iter().filter(|&&x| x).count()
283    }
284}
285
286/// Pruner for model compression
287#[derive(Debug, Clone)]
288pub struct Pruner {
289    config: PruningConfig,
290}
291
292impl Pruner {
293    /// Create a new pruner
294    pub fn new(config: PruningConfig) -> Self {
295        Self { config }
296    }
297
298    /// Generate pruning mask for weights
299    pub fn generate_mask(&self, weights: &[f32], layer_name: String) -> PruningMask {
300        let mut mask = PruningMask::new(weights.len(), layer_name.clone());
301
302        // Check if layer should be pruned
303        if self.config.layers_to_skip.contains(&layer_name) {
304            return mask; // Return mask with all weights kept
305        }
306
307        // Apply pruning strategy
308        match self.config.strategy {
309            PruningStrategy::Magnitude | PruningStrategy::L1Norm => {
310                mask.apply_magnitude_pruning(weights, self.config.sparsity);
311            }
312            PruningStrategy::Random => {
313                mask.apply_random_pruning(self.config.sparsity);
314            }
315            _ => {
316                // Other strategies can be implemented here
317                mask.apply_magnitude_pruning(weights, self.config.sparsity);
318            }
319        }
320
321        mask
322    }
323
324    /// Get configuration
325    pub fn config(&self) -> &PruningConfig {
326        &self.config
327    }
328}
329
330/// Knowledge distillation configuration
331#[derive(Debug, Clone, Serialize, Deserialize)]
332pub struct DistillationConfig {
333    /// Temperature for softening probability distributions
334    pub temperature: f32,
335    /// Weight for distillation loss (alpha)
336    /// Final loss = alpha * distillation_loss + (1 - alpha) * student_loss
337    pub alpha: f32,
338    /// Whether to use soft targets
339    pub use_soft_targets: bool,
340    /// Whether to use intermediate feature matching
341    pub use_feature_matching: bool,
342    /// Layers to match features (teacher_layer, student_layer)
343    pub feature_match_layers: Vec<(String, String)>,
344}
345
346impl DistillationConfig {
347    /// Create a new distillation configuration
348    pub fn new() -> Self {
349        Self {
350            temperature: 3.0,
351            alpha: 0.7,
352            use_soft_targets: true,
353            use_feature_matching: false,
354            feature_match_layers: Vec::new(),
355        }
356    }
357
358    /// Set temperature
359    pub fn with_temperature(mut self, temperature: f32) -> Self {
360        self.temperature = temperature.max(1.0);
361        self
362    }
363
364    /// Set alpha (distillation weight)
365    pub fn with_alpha(mut self, alpha: f32) -> Self {
366        self.alpha = alpha.clamp(0.0, 1.0);
367        self
368    }
369
370    /// Enable feature matching
371    pub fn with_feature_matching(mut self) -> Self {
372        self.use_feature_matching = true;
373        self
374    }
375
376    /// Add feature matching pair
377    pub fn add_feature_match(mut self, teacher_layer: String, student_layer: String) -> Self {
378        self.feature_match_layers
379            .push((teacher_layer, student_layer));
380        self
381    }
382}
383
384impl Default for DistillationConfig {
385    fn default() -> Self {
386        Self::new()
387    }
388}
389
390/// Knowledge distillation loss components
391#[derive(Debug, Clone, Serialize, Deserialize)]
392pub struct DistillationLoss {
393    /// Distillation loss (KL divergence between teacher and student)
394    pub distillation_loss: f32,
395    /// Student loss (cross entropy with hard labels)
396    pub student_loss: f32,
397    /// Feature matching loss (if enabled)
398    pub feature_loss: Option<f32>,
399    /// Total combined loss
400    pub total_loss: f32,
401}
402
403impl DistillationLoss {
404    /// Create a new distillation loss
405    pub fn new(distillation_loss: f32, student_loss: f32, alpha: f32) -> Self {
406        let total_loss = alpha * distillation_loss + (1.0 - alpha) * student_loss;
407
408        Self {
409            distillation_loss,
410            student_loss,
411            feature_loss: None,
412            total_loss,
413        }
414    }
415
416    /// Add feature matching loss
417    pub fn with_feature_loss(mut self, feature_loss: f32, beta: f32) -> Self {
418        self.feature_loss = Some(feature_loss);
419        self.total_loss += beta * feature_loss;
420        self
421    }
422}
423
424/// Model optimization statistics
425#[derive(Debug, Clone, Serialize, Deserialize)]
426pub struct OptimizationStats {
427    /// Original number of parameters
428    pub original_params: usize,
429    /// Number of parameters after optimization
430    pub optimized_params: usize,
431    /// Sparsity achieved
432    pub sparsity: f32,
433    /// Compression ratio
434    pub compression_ratio: f32,
435    /// Speedup estimate
436    pub speedup_estimate: f32,
437}
438
439impl OptimizationStats {
440    /// Create new optimization statistics
441    pub fn new(original_params: usize, optimized_params: usize) -> Self {
442        let sparsity = 1.0 - (optimized_params as f32 / original_params as f32);
443        let compression_ratio = original_params as f32 / optimized_params as f32;
444
445        // Estimate speedup (conservative estimate: 1.5x for 50% sparsity)
446        let speedup_estimate = 1.0 + sparsity;
447
448        Self {
449            original_params,
450            optimized_params,
451            sparsity,
452            compression_ratio,
453            speedup_estimate,
454        }
455    }
456
457    /// Parameter reduction percentage
458    pub fn param_reduction_percent(&self) -> f32 {
459        self.sparsity * 100.0
460    }
461}
462
463/// Operator fusion configuration
464#[derive(Debug, Clone, Serialize, Deserialize)]
465pub struct FusionConfig {
466    /// Fuse Conv + BatchNorm
467    pub fuse_conv_bn: bool,
468    /// Fuse Conv + ReLU
469    pub fuse_conv_relu: bool,
470    /// Fuse Linear + Bias
471    pub fuse_linear_bias: bool,
472    /// Fuse consecutive operations when possible
473    pub fuse_consecutive: bool,
474}
475
476impl FusionConfig {
477    /// Create a new fusion configuration
478    pub fn new() -> Self {
479        Self {
480            fuse_conv_bn: true,
481            fuse_conv_relu: true,
482            fuse_linear_bias: true,
483            fuse_consecutive: true,
484        }
485    }
486
487    /// Enable all fusion optimizations
488    pub fn all() -> Self {
489        Self::new()
490    }
491}
492
493impl Default for FusionConfig {
494    fn default() -> Self {
495        Self::new()
496    }
497}
498
499/// Fused operation representation
500#[derive(Debug, Clone, Serialize, Deserialize)]
501pub struct FusedOperation {
502    /// Name of the fused operation
503    pub name: String,
504    /// Original operations that were fused
505    pub operations: Vec<String>,
506    /// Estimated speedup from fusion
507    pub speedup: f32,
508}
509
510impl FusedOperation {
511    /// Create a new fused operation
512    pub fn new(name: String, operations: Vec<String>, speedup: f32) -> Self {
513        Self {
514            name,
515            operations,
516            speedup,
517        }
518    }
519}
520
521#[cfg(test)]
522mod tests {
523    use super::*;
524
525    #[test]
526    fn test_pruning_config() {
527        let config = PruningConfig::new()
528            .with_strategy(PruningStrategy::Magnitude)
529            .with_sparsity(0.7)
530            .skip_layer("output".to_string());
531
532        assert_eq!(config.strategy, PruningStrategy::Magnitude);
533        assert_eq!(config.sparsity, 0.7);
534        assert!(config.layers_to_skip.contains(&"output".to_string()));
535    }
536
537    #[test]
538    fn test_pruning_mask_creation() {
539        let mask = PruningMask::new(100, "layer1".to_string());
540
541        assert_eq!(mask.mask.len(), 100);
542        assert_eq!(mask.sparsity, 0.0); // No pruning initially
543        assert_eq!(mask.num_kept(), 100);
544        assert_eq!(mask.num_pruned(), 0);
545    }
546
547    #[test]
548    fn test_magnitude_pruning() {
549        let weights = vec![0.1, 0.5, 0.2, 0.8, 0.3];
550        let mut mask = PruningMask::new(weights.len(), "test".to_string());
551
552        mask.apply_magnitude_pruning(&weights, 0.4); // Prune 40%
553
554        assert_eq!(mask.num_pruned(), 2); // Should prune 2 out of 5
555        assert!((mask.sparsity - 0.4).abs() < 0.1);
556    }
557
558    #[test]
559    fn test_random_pruning() {
560        let mut mask = PruningMask::new(100, "test".to_string());
561
562        mask.apply_random_pruning(0.5);
563
564        // Random pruning may select same index multiple times, so count may vary
565        // Should prune at least 30% but likely close to 50% (with some overlap)
566        let pruned = mask.num_pruned();
567        assert!(
568            pruned >= 30 && pruned <= 60,
569            "Expected ~50 pruned, got {}",
570            pruned
571        );
572        assert!(mask.sparsity > 0.2 && mask.sparsity < 0.7);
573    }
574
575    #[test]
576    fn test_mask_application() {
577        let mut weights = vec![1.0, 2.0, 3.0, 4.0, 5.0];
578        let mut mask = PruningMask::new(weights.len(), "test".to_string());
579
580        // Manually prune first two weights
581        mask.mask[0] = false;
582        mask.mask[1] = false;
583        mask.update_sparsity();
584
585        mask.apply_to_weights(&mut weights);
586
587        assert_eq!(weights[0], 0.0);
588        assert_eq!(weights[1], 0.0);
589        assert_eq!(weights[2], 3.0);
590        assert_eq!(weights[3], 4.0);
591        assert_eq!(weights[4], 5.0);
592    }
593
594    #[test]
595    fn test_pruner_generation() {
596        let config = PruningConfig::new().with_sparsity(0.5);
597        let pruner = Pruner::new(config);
598
599        let weights = vec![0.1, 0.5, 0.2, 0.8, 0.3, 0.9, 0.1, 0.4];
600        let mask = pruner.generate_mask(&weights, "layer1".to_string());
601
602        assert_eq!(mask.num_pruned(), 4); // 50% of 8
603        assert!((mask.sparsity - 0.5).abs() < 0.01);
604    }
605
606    #[test]
607    fn test_pruner_skip_layer() {
608        let config = PruningConfig::new()
609            .with_sparsity(0.5)
610            .skip_layer("output".to_string());
611
612        let pruner = Pruner::new(config);
613
614        let weights = vec![0.1, 0.5, 0.2, 0.8];
615        let mask = pruner.generate_mask(&weights, "output".to_string());
616
617        assert_eq!(mask.num_pruned(), 0); // Should not prune
618        assert_eq!(mask.sparsity, 0.0);
619    }
620
621    #[test]
622    fn test_distillation_config() {
623        let config = DistillationConfig::new()
624            .with_temperature(4.0)
625            .with_alpha(0.8)
626            .with_feature_matching();
627
628        assert_eq!(config.temperature, 4.0);
629        assert_eq!(config.alpha, 0.8);
630        assert!(config.use_feature_matching);
631    }
632
633    #[test]
634    fn test_distillation_loss() {
635        let loss = DistillationLoss::new(0.5, 0.3, 0.7);
636
637        assert_eq!(loss.distillation_loss, 0.5);
638        assert_eq!(loss.student_loss, 0.3);
639        assert!((loss.total_loss - (0.7 * 0.5 + 0.3 * 0.3)).abs() < 0.001);
640    }
641
642    #[test]
643    fn test_distillation_loss_with_feature() {
644        let loss = DistillationLoss::new(0.5, 0.3, 0.7).with_feature_loss(0.2, 0.1);
645
646        assert!(loss.feature_loss.is_some());
647        assert_eq!(loss.feature_loss.unwrap(), 0.2);
648        assert!(loss.total_loss > 0.4); // Should be higher with feature loss
649    }
650
651    #[test]
652    fn test_optimization_stats() {
653        let stats = OptimizationStats::new(1000, 500);
654
655        assert_eq!(stats.original_params, 1000);
656        assert_eq!(stats.optimized_params, 500);
657        assert_eq!(stats.sparsity, 0.5);
658        assert_eq!(stats.compression_ratio, 2.0);
659        assert_eq!(stats.param_reduction_percent(), 50.0);
660    }
661
662    #[test]
663    fn test_fusion_config() {
664        let config = FusionConfig::new();
665
666        assert!(config.fuse_conv_bn);
667        assert!(config.fuse_conv_relu);
668        assert!(config.fuse_linear_bias);
669    }
670
671    #[test]
672    fn test_fused_operation() {
673        let fused = FusedOperation::new(
674            "conv_bn_relu".to_string(),
675            vec!["conv".to_string(), "bn".to_string(), "relu".to_string()],
676            1.5,
677        );
678
679        assert_eq!(fused.name, "conv_bn_relu");
680        assert_eq!(fused.operations.len(), 3);
681        assert_eq!(fused.speedup, 1.5);
682    }
683
684    #[test]
685    fn test_pruning_schedule() {
686        let one_shot = PruningSchedule::OneShot;
687        let gradual = PruningSchedule::Gradual { steps: 1000 };
688        let iterative = PruningSchedule::Iterative {
689            iterations: 5,
690            train_steps: 100,
691        };
692
693        assert_eq!(one_shot, PruningSchedule::OneShot);
694        assert_ne!(gradual, one_shot);
695        assert_ne!(iterative, gradual);
696    }
697}
torsh_ffi/model_optimization.rs

torsh_ffi/
model_optimization.rs
