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
41impl ModuleMapper for Quantizer {
42    fn enter_module(&mut self, name: &str, _container_type: &str) {
43        self.path.push(name.to_string());
44    }
45
46    fn exit_module(&mut self, _name: &str, _container_type: &str) {
47        self.path.pop();
48    }
49
50    fn map_float<const D: usize>(&mut self, param: Param<Tensor<D>>) -> Param<Tensor<D>> {
51        let (id, mut tensor, mapper) = param.consume();
52        let path = self.path.join(".");
53        if self.group.matches(&id, Some(&path)) {
54            let range = compute_range(&self.scheme, &tensor, &self.calibration);
55            let qparams = compute_q_params(&self.scheme, range);
56            tensor = tensor.quantize(&self.scheme, qparams);
57        }
58        Param::from_mapped_value(id, tensor, mapper)
59    }
60}
61
62#[cfg(all(test, not(feature = "tch")))]
63mod tests {
64    use crate::module::{Module, ParamGroup, Quantizer};
65    use crate::tensor::DType;
66    use crate::test_device;
67    use crate::test_utils::SimpleLinear;
68    use burn_tensor::{
69        Device, Tensor, Tolerance,
70        quantization::{Calibration, QuantLevel, QuantParam, QuantScheme, QuantValue},
71    };
72
73    /// Per-tensor Q8 symmetric scheme used across these tests.
74    fn test_scheme(device: &Device) -> QuantScheme {
75        device
76            .settings()
77            .quantization
78            .scheme
79            .with_value(QuantValue::Q8S)
80            .with_level(QuantLevel::Tensor)
81            .with_param(QuantParam::F32)
82    }
83
84    /// Whether a tensor currently holds quantized (`QFloat`) data.
85    fn is_quantized<const D: usize>(tensor: &Tensor<D>) -> bool {
86        matches!(tensor.dtype(), DType::QFloat(_))
87    }
88
89    #[test]
90    fn should_quantize_module() {
91        let device = test_device();
92        let module = SimpleLinear::new(32, 32, &device);
93        let scheme = test_scheme(&device);
94
95        let result = module.weight.val();
96
97        let calibration = Calibration::MinMax;
98        let mut quantizer = Quantizer::new(calibration, scheme);
99        let q_module = module.quantize_weights(&mut quantizer);
100        let q_result = q_module.weight.val().dequantize();
101
102        result
103            .into_data()
104            .assert_approx_eq::<f32>(&q_result.into_data(), Tolerance::permissive());
105    }
106
107    #[test]
108    fn should_quantize_only_group() {
109        let device = test_device();
110        let module = SimpleLinear::new(32, 32, &device);
111        let scheme = test_scheme(&device);
112
113        let weight_ref = module.weight.val();
114
115        let mut quantizer = Quantizer::new(Calibration::MinMax, scheme);
116        let q_module =
117            module.quantize_weights_group(&mut quantizer, ParamGroup::from_path("weight"));
118
119        // Only the weight (the group) is quantized; the bias is left untouched.
120        assert!(is_quantized(&q_module.weight.val()));
121        let bias = q_module.bias.clone().expect("bias should be present");
122        assert!(!is_quantized(&bias.val()));
123
124        // The quantized weight still approximates the original values.
125        weight_ref.into_data().assert_approx_eq::<f32>(
126            &q_module.weight.val().dequantize().into_data(),
127            Tolerance::permissive(),
128        );
129    }
130
131    #[test]
132    fn should_quantize_all_when_no_group() {
133        let device = test_device();
134        let module = SimpleLinear::new(32, 32, &device);
135        let scheme = test_scheme(&device);
136
137        let mut quantizer = Quantizer::new(Calibration::MinMax, scheme);
138        let q_module = module.quantize_weights(&mut quantizer);
139
140        // Without a group every float parameter is quantized.
141        assert!(is_quantized(&q_module.weight.val()));
142        assert!(is_quantized(&q_module.bias.clone().unwrap().val()));
143    }
144}
burn_core/module/quantize.rs

burn_core/module/
quantize.rs
