burn_core/module/
quantize.rs

1use alloc::{
2    string::{String, ToString},
3    vec::Vec,
4};
5
6use burn_tensor::{
7    Tensor,
8    quantization::{Calibration, QuantScheme, compute_q_params, compute_range},
9};
10
11use crate::module::{ModuleMapper, Param, ParamGroup};
12
13/// Describes how to quantize a module.
14pub struct Quantizer {
15    path: Vec<String>,
16    /// The calibration method used in quantization.
17    pub calibration: Calibration,
18    /// The quantization scheme.
19    pub scheme: QuantScheme,
20    /// The parameter group to quantize.
21    pub group: ParamGroup,
22}
23
24impl Quantizer {
25    /// Create a new [Quantizer].
26    pub fn new(calibration: Calibration, scheme: QuantScheme) -> Self {
27        Self {
28            path: alloc::vec![],
29            calibration,
30            scheme,
31            group: ParamGroup::all(),
32        }
33    }
34
35    /// Set the parameter group to quantize.
36    pub fn set_param_group(&mut self, group: ParamGroup) {
37        self.group = group
38    }
39
40    pub(crate) fn map_float_at_path<const D: usize>(
41        &self,
42        param: Param<Tensor<D>>,
43        path: &str,
44    ) -> Param<Tensor<D>> {
45        if self.group.matches(&param.id, Some(path)) {
46            return param.map(|tensor| {
47                let range = compute_range(&self.scheme, &tensor, &self.calibration);
48                let qparams = compute_q_params(&self.scheme, range);
49                tensor.quantize(&self.scheme, qparams)
50            });
51        }
52
53        param
54    }
55}
56
57impl ModuleMapper for Quantizer {
58    fn enter_module(&mut self, name: &str, _container_type: &str) {
59        self.path.push(name.to_string());
60    }
61
62    fn exit_module(&mut self, _name: &str, _container_type: &str) {
63        self.path.pop();
64    }
65
66    fn map_float<const D: usize>(&mut self, param: Param<Tensor<D>>) -> Param<Tensor<D>> {
67        let path = self.path.join(".");
68        self.map_float_at_path(param, &path)
69    }
70}
71
72#[cfg(all(test, not(feature = "tch")))]
73mod tests {
74    use crate::module::{Module, ParamGroup, Quantizer};
75    use crate::tensor::DType;
76    use crate::test_device;
77    use crate::test_utils::SimpleLinear;
78    use burn_tensor::{
79        Device, Tensor, Tolerance,
80        quantization::{Calibration, QuantScheme, QuantValue},
81    };
82
83    /// Per-tensor Q8 symmetric scheme used across these tests.
84    fn test_scheme(device: &Device) -> QuantScheme {
85        device
86            .settings()
87            .quantization
88            .scheme
89            .with_value(QuantValue::Q8S)
90    }
91
92    /// Whether a tensor currently holds quantized (`QFloat`) data.
93    fn is_quantized<const D: usize>(tensor: &Tensor<D>) -> bool {
94        matches!(tensor.dtype(), DType::QFloat(_))
95    }
96
97    #[test]
98    fn should_quantize_module() {
99        let device = test_device();
100        let module = SimpleLinear::new(32, 32, &device);
101        let scheme = test_scheme(&device);
102
103        let result = module.weight.val();
104
105        let calibration = Calibration::MinMax;
106        let mut quantizer = Quantizer::new(calibration, scheme);
107        let q_module = module.quantize_weights(&mut quantizer);
108        let q_result = q_module.weight.val().dequantize();
109
110        result
111            .into_data()
112            .assert_approx_eq::<f32>(&q_result.into_data(), Tolerance::permissive());
113    }
114
115    #[test]
116    fn should_quantize_only_group() {
117        let device = test_device();
118        let module = SimpleLinear::new(32, 32, &device);
119        let scheme = test_scheme(&device);
120
121        let weight_ref = module.weight.val();
122
123        let mut quantizer = Quantizer::new(Calibration::MinMax, scheme);
124        let q_module =
125            module.quantize_weights_group(&mut quantizer, ParamGroup::from_path("weight"));
126
127        // Only the weight (the group) is quantized; the bias is left untouched.
128        assert!(is_quantized(&q_module.weight.val()));
129        let bias = q_module.bias.clone().expect("bias should be present");
130        assert!(!is_quantized(&bias.val()));
131
132        // The quantized weight still approximates the original values.
133        weight_ref.into_data().assert_approx_eq::<f32>(
134            &q_module.weight.val().dequantize().into_data(),
135            Tolerance::permissive(),
136        );
137    }
138
139    #[test]
140    fn should_quantize_all_when_no_group() {
141        let device = test_device();
142        let module = SimpleLinear::new(32, 32, &device);
143        let scheme = test_scheme(&device);
144
145        let mut quantizer = Quantizer::new(Calibration::MinMax, scheme);
146        let q_module = module.quantize_weights(&mut quantizer);
147
148        // Without a group every float parameter is quantized.
149        assert!(is_quantized(&q_module.weight.val()));
150        assert!(is_quantized(&q_module.bias.clone().unwrap().val()));
151    }
152}
burn_core/module/quantize.rs

burn_core/module/
quantize.rs
