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        let (id, mut tensor, mapper) = param.consume();
46        if self.group.matches(&id, Some(path)) {
47            let range = compute_range(&self.scheme, &tensor, &self.calibration);
48            let qparams = compute_q_params(&self.scheme, range);
49            tensor = tensor.quantize(&self.scheme, qparams);
50        }
51        Param::from_mapped_value(id, tensor, mapper)
52    }
53}
54
55impl ModuleMapper for Quantizer {
56    fn enter_module(&mut self, name: &str, _container_type: &str) {
57        self.path.push(name.to_string());
58    }
59
60    fn exit_module(&mut self, _name: &str, _container_type: &str) {
61        self.path.pop();
62    }
63
64    fn map_float<const D: usize>(&mut self, param: Param<Tensor<D>>) -> Param<Tensor<D>> {
65        let path = self.path.join(".");
66        self.map_float_at_path(param, &path)
67    }
68}
69
70#[cfg(all(test, not(feature = "tch")))]
71mod tests {
72    use crate::module::{Module, ParamGroup, Quantizer};
73    use crate::tensor::DType;
74    use crate::test_device;
75    use crate::test_utils::SimpleLinear;
76    use burn_tensor::{
77        Device, Tensor, Tolerance,
78        quantization::{Calibration, QuantLevel, QuantParam, QuantScheme, QuantValue},
79    };
80
81    /// Per-tensor Q8 symmetric scheme used across these tests.
82    fn test_scheme(device: &Device) -> QuantScheme {
83        device
84            .settings()
85            .quantization
86            .scheme
87            .with_value(QuantValue::Q8S)
88            .with_level(QuantLevel::Tensor)
89            .with_param(QuantParam::F32)
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
