scirs2_io/ml_framework/
quantization.rs

1//! Model quantization support
2
3use crate::error::{IoError, Result};
4use crate::ml_framework::{MLModel, MLTensor, ModelMetadata, TensorMetadata};
5use scirs2_core::ndarray::{ArrayD, IxDyn};
6use std::collections::HashMap;
7
8/// Quantization methods
9#[derive(Debug, Clone, Copy)]
10pub enum QuantizationMethod {
11    /// Dynamic quantization
12    Dynamic,
13    /// Static quantization with calibration
14    Static,
15    /// Quantization-aware training
16    QAT,
17    /// Post-training quantization
18    PTQ,
19}
20
21/// Quantized tensor
22#[derive(Debug, Clone)]
23pub struct QuantizedTensor {
24    pub data: Vec<u8>,
25    pub scale: f32,
26    pub zero_point: i32,
27    pub metadata: TensorMetadata,
28}
29
30impl QuantizedTensor {
31    /// Quantize a floating-point tensor
32    pub fn from_float_tensor(tensor: &MLTensor, bits: u8) -> Result<Self> {
33        let data = tensor.data.as_slice().expect("Operation failed");
34        let (min_val, max_val) = data
35            .iter()
36            .fold((f32::INFINITY, f32::NEG_INFINITY), |(min, max), &x| {
37                (min.min(x), max.max(x))
38            });
39
40        let qmax = (1 << bits) - 1;
41        let scale = (max_val - min_val) / qmax as f32;
42        let zero_point = (-min_val / scale).round() as i32;
43
44        let quantized: Vec<u8> = data
45            .iter()
46            .map(|&x| (x / scale + zero_point as f32).round() as u8)
47            .collect();
48
49        Ok(Self {
50            data: quantized,
51            scale,
52            zero_point,
53            metadata: tensor.metadata.clone(),
54        })
55    }
56
57    /// Dequantize to floating-point
58    pub fn to_float_tensor(&self) -> Result<MLTensor> {
59        let data: Vec<f32> = self
60            .data
61            .iter()
62            .map(|&q| (q as i32 - self.zero_point) as f32 * self.scale)
63            .collect();
64
65        let array = ArrayD::from_shape_vec(IxDyn(&self.metadata.shape), data)
66            .map_err(|e| IoError::Other(e.to_string()))?;
67
68        Ok(MLTensor::new(array, self.metadata.name.clone()))
69    }
70}
71
72/// Model quantizer
73pub struct ModelQuantizer {
74    method: QuantizationMethod,
75    bits: u8,
76}
77
78impl ModelQuantizer {
79    pub fn new(method: QuantizationMethod, bits: u8) -> Self {
80        Self { method, bits }
81    }
82
83    /// Quantize entire model
84    pub fn quantize_model(&self, model: &MLModel) -> Result<QuantizedModel> {
85        let mut quantized_weights = HashMap::new();
86
87        for (name, tensor) in &model.weights {
88            let quantized = QuantizedTensor::from_float_tensor(tensor, self.bits)?;
89            quantized_weights.insert(name.clone(), quantized);
90        }
91
92        Ok(QuantizedModel {
93            metadata: model.metadata.clone(),
94            weights: quantized_weights,
95            config: model.config.clone(),
96            quantization_info: QuantizationInfo {
97                method: self.method,
98                bits: self.bits,
99            },
100        })
101    }
102}
103
104/// Quantized model
105#[derive(Debug, Clone)]
106pub struct QuantizedModel {
107    pub metadata: ModelMetadata,
108    pub weights: HashMap<String, QuantizedTensor>,
109    pub config: HashMap<String, serde_json::Value>,
110    pub quantization_info: QuantizationInfo,
111}
112
113#[derive(Debug, Clone)]
114pub struct QuantizationInfo {
115    pub method: QuantizationMethod,
116    pub bits: u8,
117}
scirs2_io/ml_framework/quantization.rs

scirs2_io/ml_framework/
quantization.rs
