torsh_hub/
quantization.rs

1//! # Model Quantization and Compression
2//!
3//! This module provides comprehensive quantization utilities for reducing model size
4//! and improving inference performance while maintaining accuracy.
5//!
6//! ## Features
7//!
8//! - **Post-Training Quantization**: Quantize pretrained models without retraining
9//! - **Quantization-Aware Training**: Fine-tune models with quantization simulation
10//! - **Multiple Precision Levels**: INT8, FP16, and mixed precision support
11//! - **Calibration**: Automatic calibration using representative data
12//! - **Layer-wise Quantization**: Fine-grained control over which layers to quantize
13//!
14//! ## SciRS2 POLICY Compliance
15//!
16//! This module strictly follows the SciRS2 POLICY:
17//! - Uses `scirs2_core::ndarray::*` for all array operations
18//! - Uses `scirs2_core::random::*` for calibration sampling
19//!
20//! ## Quick Start
21//!
22//! ```rust,no_run
23//! use torsh_hub::quantization::{QuantizationConfig, ModelQuantizer, QuantizationType};
24//! use torsh_nn::Module;
25//!
26//! # fn example(model: Box<dyn Module>) -> Result<(), Box<dyn std::error::Error>> {
27//! // Configure quantization
28//! let config = QuantizationConfig {
29//!     quantization_type: QuantizationType::Int8,
30//!     per_channel: true,
31//!     symmetric: true,
32//!     calibration_samples: 100,
33//!     excluded_layers: vec![],
34//! };
35//!
36//! // Quantize the model
37//! let mut quantizer = ModelQuantizer::new(config);
38//! let quantized_model = quantizer.quantize(model)?;
39//!
40//! // Model is now optimized for deployment
41//! # Ok(())
42//! # }
43//! ```
44
45use serde::{Deserialize, Serialize};
46use std::collections::HashMap;
47use torsh_core::error::{Result, TorshError};
48
49// SciRS2 POLICY: Use unified ndarray access
50use scirs2_core::ndarray::{Array1, Array2, ArrayD};
51
52/// Quantization configuration
53#[derive(Debug, Clone, Serialize, Deserialize)]
54pub struct QuantizationConfig {
55    /// Type of quantization to apply
56    pub quantization_type: QuantizationType,
57    /// Use per-channel quantization (better accuracy, larger overhead)
58    pub per_channel: bool,
59    /// Use symmetric quantization (centered at zero)
60    pub symmetric: bool,
61    /// Number of calibration samples for activation quantization
62    pub calibration_samples: usize,
63    /// Layers to exclude from quantization
64    pub excluded_layers: Vec<String>,
65}
66
67impl Default for QuantizationConfig {
68    fn default() -> Self {
69        Self {
70            quantization_type: QuantizationType::Int8,
71            per_channel: true,
72            symmetric: true,
73            calibration_samples: 100,
74            excluded_layers: vec![],
75        }
76    }
77}
78
79/// Supported quantization types
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
81pub enum QuantizationType {
82    /// 8-bit integer quantization (INT8)
83    Int8,
84    /// 16-bit floating point (FP16/half precision)
85    Fp16,
86    /// Mixed precision (critical layers in FP16, others in INT8)
87    Mixed,
88    /// Dynamic quantization (weights quantized, activations in FP32)
89    Dynamic,
90}
91
92/// Model quantizer for post-training quantization
93pub struct ModelQuantizer {
94    config: QuantizationConfig,
95    calibration_stats: HashMap<String, QuantizationStats>,
96}
97
98impl ModelQuantizer {
99    /// Create a new model quantizer
100    ///
101    /// # Arguments
102    /// * `config` - Quantization configuration
103    ///
104    /// # Example
105    /// ```rust
106    /// use torsh_hub::quantization::{ModelQuantizer, QuantizationConfig};
107    ///
108    /// let config = QuantizationConfig::default();
109    /// let quantizer = ModelQuantizer::new(config);
110    /// ```
111    pub fn new(config: QuantizationConfig) -> Self {
112        Self {
113            config,
114            calibration_stats: HashMap::new(),
115        }
116    }
117
118    /// Quantize a model using the configured settings
119    ///
120    /// This performs post-training quantization on the model.
121    ///
122    /// # Arguments
123    /// * `model` - The model to quantize
124    ///
125    /// # Returns
126    /// * Quantized model wrapped in Result
127    pub fn quantize(
128        &mut self,
129        _model: Box<dyn torsh_nn::Module>,
130    ) -> Result<Box<dyn torsh_nn::Module>> {
131        // Implementation would convert model weights and potentially activations
132        // to lower precision based on config
133        Err(TorshError::NotImplemented(
134            "Full model quantization requires torsh-nn integration".to_string(),
135        ))
136    }
137
138    /// Calibrate quantization parameters using sample data
139    ///
140    /// # Arguments
141    /// * `layer_name` - Name of the layer to calibrate
142    /// * `activations` - Sample activations for calibration
143    pub fn calibrate(&mut self, layer_name: &str, activations: &ArrayD<f32>) -> Result<()> {
144        let stats = match self.config.quantization_type {
145            QuantizationType::Int8 => self.compute_int8_stats(activations)?,
146            QuantizationType::Fp16 => self.compute_fp16_stats(activations)?,
147            QuantizationType::Mixed => self.compute_mixed_stats(activations)?,
148            QuantizationType::Dynamic => self.compute_dynamic_stats(activations)?,
149        };
150
151        self.calibration_stats.insert(layer_name.to_string(), stats);
152        Ok(())
153    }
154
155    /// Compute INT8 quantization statistics
156    fn compute_int8_stats(&self, activations: &ArrayD<f32>) -> Result<QuantizationStats> {
157        let min = activations
158            .iter()
159            .copied()
160            .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
161            .unwrap_or(0.0);
162        let max = activations
163            .iter()
164            .copied()
165            .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
166            .unwrap_or(0.0);
167
168        let (scale, zero_point) = if self.config.symmetric {
169            // Symmetric: map [-abs_max, abs_max] to [-127, 127]
170            let abs_max = min.abs().max(max.abs());
171            let scale = abs_max / 127.0;
172            (scale, 0)
173        } else {
174            // Asymmetric: map [min, max] to [0, 255]
175            let scale = (max - min) / 255.0;
176            let zero_point = (-min / scale).round() as i32;
177            (scale, zero_point)
178        };
179
180        Ok(QuantizationStats {
181            scale,
182            zero_point,
183            min,
184            max,
185            num_bits: 8,
186        })
187    }
188
189    /// Compute FP16 quantization statistics
190    fn compute_fp16_stats(&self, _activations: &ArrayD<f32>) -> Result<QuantizationStats> {
191        // FP16 doesn't need calibration stats, but we store metadata
192        Ok(QuantizationStats {
193            scale: 1.0,
194            zero_point: 0,
195            min: 0.0,
196            max: 0.0,
197            num_bits: 16,
198        })
199    }
200
201    /// Compute mixed precision statistics
202    fn compute_mixed_stats(&self, activations: &ArrayD<f32>) -> Result<QuantizationStats> {
203        // For mixed precision, use INT8 stats but mark critical layers
204        self.compute_int8_stats(activations)
205    }
206
207    /// Compute dynamic quantization statistics
208    fn compute_dynamic_stats(&self, activations: &ArrayD<f32>) -> Result<QuantizationStats> {
209        // Dynamic quantization computes stats on the fly during inference
210        self.compute_int8_stats(activations)
211    }
212
213    /// Get calibration statistics for a layer
214    pub fn get_stats(&self, layer_name: &str) -> Option<&QuantizationStats> {
215        self.calibration_stats.get(layer_name)
216    }
217
218    /// Export quantization configuration
219    pub fn export_config(&self) -> Result<String> {
220        serde_json::to_string_pretty(&self.config)
221            .map_err(|e| TorshError::SerializationError(e.to_string()))
222    }
223}
224
225/// Quantization statistics for a layer
226#[derive(Debug, Clone, Serialize, Deserialize)]
227pub struct QuantizationStats {
228    /// Scaling factor for quantization
229    pub scale: f32,
230    /// Zero point for asymmetric quantization
231    pub zero_point: i32,
232    /// Minimum activation value observed
233    pub min: f32,
234    /// Maximum activation value observed
235    pub max: f32,
236    /// Number of bits used
237    pub num_bits: u8,
238}
239
240/// Quantize a weight tensor to INT8
241///
242/// # Arguments
243/// * `weights` - FP32 weights to quantize
244/// * `per_channel` - Whether to use per-channel quantization
245///
246/// # Returns
247/// * Tuple of (quantized weights, scale factors, zero points)
248///
249/// # Example
250/// ```rust
251/// use torsh_hub::quantization::quantize_weights_int8;
252/// use scirs2_core::ndarray::Array2;
253///
254/// let weights = Array2::<f32>::zeros((64, 128));
255/// let (quantized, scales, zeros) = quantize_weights_int8(&weights, true);
256/// ```
257pub fn quantize_weights_int8(
258    weights: &Array2<f32>,
259    per_channel: bool,
260) -> (Array2<i8>, Array1<f32>, Array1<i32>) {
261    let (rows, cols) = weights.dim();
262
263    if per_channel {
264        // Per-channel quantization (separate scale per output channel)
265        let mut quantized = Array2::zeros((rows, cols));
266        let mut scales = Array1::zeros(rows);
267        let mut zero_points = Array1::zeros(rows);
268
269        for i in 0..rows {
270            let row = weights.row(i);
271            let min = row
272                .iter()
273                .copied()
274                .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
275                .unwrap_or(0.0);
276            let max = row
277                .iter()
278                .copied()
279                .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
280                .unwrap_or(0.0);
281
282            let scale = (max - min) / 255.0;
283            let zero_point = if scale > 0.0 {
284                (-min / scale).round() as i32
285            } else {
286                0
287            };
288
289            scales[i] = scale;
290            zero_points[i] = zero_point;
291
292            for j in 0..cols {
293                let val = if scale > 0.0 {
294                    ((weights[[i, j]] / scale) + zero_point as f32)
295                        .round()
296                        .clamp(-128.0, 127.0) as i8
297                } else {
298                    0
299                };
300                quantized[[i, j]] = val;
301            }
302        }
303
304        (quantized, scales, zero_points)
305    } else {
306        // Per-tensor quantization (single scale for entire tensor)
307        let min = weights
308            .iter()
309            .copied()
310            .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
311            .unwrap_or(0.0);
312        let max = weights
313            .iter()
314            .copied()
315            .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
316            .unwrap_or(0.0);
317
318        let scale = (max - min) / 255.0;
319        let zero_point = if scale > 0.0 {
320            (-min / scale).round() as i32
321        } else {
322            0
323        };
324
325        let quantized = weights.mapv(|w| {
326            if scale > 0.0 {
327                ((w / scale) + zero_point as f32)
328                    .round()
329                    .clamp(-128.0, 127.0) as i8
330            } else {
331                0
332            }
333        });
334
335        let scales = Array1::from_elem(1, scale);
336        let zero_points = Array1::from_elem(1, zero_point);
337
338        (quantized, scales, zero_points)
339    }
340}
341
342/// Dequantize INT8 weights back to FP32
343///
344/// # Arguments
345/// * `quantized` - Quantized INT8 weights
346/// * `scales` - Scale factors
347/// * `zero_points` - Zero points
348///
349/// # Returns
350/// * Dequantized FP32 weights
351pub fn dequantize_weights_int8(
352    quantized: &Array2<i8>,
353    scales: &Array1<f32>,
354    zero_points: &Array1<i32>,
355) -> Array2<f32> {
356    let (rows, cols) = quantized.dim();
357    let per_channel = scales.len() > 1;
358
359    if per_channel {
360        let mut dequantized = Array2::zeros((rows, cols));
361        for i in 0..rows {
362            let scale = scales[i];
363            let zero_point = zero_points[i];
364            for j in 0..cols {
365                dequantized[[i, j]] = (quantized[[i, j]] as f32 - zero_point as f32) * scale;
366            }
367        }
368        dequantized
369    } else {
370        let scale = scales[0];
371        let zero_point = zero_points[0];
372        quantized.mapv(|q| (q as f32 - zero_point as f32) * scale)
373    }
374}
375
376/// Calculate quantization error metrics
377///
378/// # Arguments
379/// * `original` - Original FP32 weights
380/// * `quantized` - Quantized weights
381/// * `scales` - Scale factors
382/// * `zero_points` - Zero points
383///
384/// # Returns
385/// * Tuple of (MSE, max_error, SQNR in dB)
386pub fn quantization_metrics(
387    original: &Array2<f32>,
388    quantized: &Array2<i8>,
389    scales: &Array1<f32>,
390    zero_points: &Array1<i32>,
391) -> (f32, f32, f32) {
392    let dequantized = dequantize_weights_int8(quantized, scales, zero_points);
393
394    // Mean Squared Error
395    let diff = original - &dequantized;
396    let mse = diff.mapv(|x| x * x).mean().unwrap_or(0.0);
397
398    // Maximum absolute error
399    let max_error = diff
400        .iter()
401        .map(|x| x.abs())
402        .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
403        .unwrap_or(0.0);
404
405    // Signal-to-Quantization-Noise Ratio (SQNR) in dB
406    let signal_power = original.mapv(|x| x * x).mean().unwrap_or(1e-10);
407    let sqnr_db = 10.0 * (signal_power / mse.max(1e-10)).log10();
408
409    (mse, max_error, sqnr_db)
410}
411
412#[cfg(test)]
413mod tests {
414    use super::*;
415
416    #[test]
417    fn test_quantization_config_default() {
418        let config = QuantizationConfig::default();
419        assert_eq!(config.quantization_type, QuantizationType::Int8);
420        assert!(config.per_channel);
421        assert!(config.symmetric);
422        assert_eq!(config.calibration_samples, 100);
423    }
424
425    #[test]
426    fn test_model_quantizer_creation() {
427        let config = QuantizationConfig::default();
428        let quantizer = ModelQuantizer::new(config);
429        assert_eq!(quantizer.calibration_stats.len(), 0);
430    }
431
432    #[test]
433    fn test_quantize_weights_per_tensor() {
434        let weights = Array2::from_shape_fn((4, 4), |(i, j)| (i * 4 + j) as f32);
435        let (quantized, scales, zeros) = quantize_weights_int8(&weights, false);
436
437        assert_eq!(quantized.dim(), weights.dim());
438        assert_eq!(scales.len(), 1);
439        assert_eq!(zeros.len(), 1);
440    }
441
442    #[test]
443    fn test_quantize_weights_per_channel() {
444        let weights = Array2::from_shape_fn((4, 4), |(i, j)| (i * 4 + j) as f32);
445        let (quantized, scales, zeros) = quantize_weights_int8(&weights, true);
446
447        assert_eq!(quantized.dim(), weights.dim());
448        assert_eq!(scales.len(), 4); // One per row/channel
449        assert_eq!(zeros.len(), 4);
450    }
451
452    #[test]
453    fn test_dequantize_weights() {
454        let weights = Array2::from_shape_fn((4, 4), |(i, j)| ((i * 4 + j) as f32) * 0.1);
455        let (quantized, scales, zeros) = quantize_weights_int8(&weights, false);
456        let dequantized = dequantize_weights_int8(&quantized, &scales, &zeros);
457
458        // Check that dequantization is close to original
459        let diff = &weights - &dequantized;
460        let max_diff = diff
461            .iter()
462            .map(|x| x.abs())
463            .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
464            .unwrap_or(0.0);
465
466        // Quantization error should be reasonable for INT8 quantization
467        // INT8 has 256 levels, so error can be up to scale/2 per value
468        assert!(max_diff < 1.0, "Max diff: {}", max_diff);
469    }
470
471    #[test]
472    fn test_quantization_metrics() {
473        let original = Array2::from_shape_fn((8, 8), |(i, j)| (i * 8 + j) as f32 * 0.01);
474        let (quantized, scales, zeros) = quantize_weights_int8(&original, true);
475        let (mse, max_error, sqnr) = quantization_metrics(&original, &quantized, &scales, &zeros);
476
477        // Basic sanity checks
478        assert!(mse >= 0.0);
479        assert!(max_error >= 0.0);
480        assert!(sqnr.is_finite());
481    }
482
483    #[test]
484    fn test_calibration_int8() {
485        let mut quantizer = ModelQuantizer::new(QuantizationConfig::default());
486        let activations =
487            ArrayD::from_shape_vec(vec![2, 3, 4], (0..24).map(|x| x as f32 * 0.1).collect())
488                .unwrap();
489
490        quantizer.calibrate("layer1", &activations).unwrap();
491        let stats = quantizer.get_stats("layer1").unwrap();
492
493        assert_eq!(stats.num_bits, 8);
494        assert!(stats.scale > 0.0);
495    }
496
497    #[test]
498    fn test_export_config() {
499        let quantizer = ModelQuantizer::new(QuantizationConfig::default());
500        let config_json = quantizer.export_config().unwrap();
501
502        assert!(config_json.contains("Int8"));
503        assert!(config_json.contains("per_channel"));
504    }
505}
torsh_hub/quantization.rs

torsh_hub/
quantization.rs
