burn-core 0.22.0-pre.1

Flexible and Comprehensive Deep Learning Framework in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

use alloc::{
    string::{String, ToString},
    vec::Vec,
};

use burn_tensor::{
    Tensor,
    quantization::{Calibration, QuantScheme, compute_q_params, compute_range},
};

use crate::module::{ModuleMapper, Param, ParamGroup};

/// Describes how to quantize a module.
pub struct Quantizer {
    path: Vec<String>,
    /// The calibration method used in quantization.
    pub calibration: Calibration,
    /// The quantization scheme.
    pub scheme: QuantScheme,
    /// The parameter group to quantize.
    pub group: ParamGroup,
}

impl Quantizer {
    /// Create a new [Quantizer].
    pub fn new(calibration: Calibration, scheme: QuantScheme) -> Self {
        Self {
            path: alloc::vec![],
            calibration,
            scheme,
            group: ParamGroup::all(),
        }
    }

    /// Set the parameter group to quantize.
    pub fn set_param_group(&mut self, group: ParamGroup) {
        self.group = group
    }
}

impl ModuleMapper for Quantizer {
    fn enter_module(&mut self, name: &str, _container_type: &str) {
        self.path.push(name.to_string());
    }

    fn exit_module(&mut self, _name: &str, _container_type: &str) {
        self.path.pop();
    }

    fn map_float<const D: usize>(&mut self, param: Param<Tensor<D>>) -> Param<Tensor<D>> {
        let (id, mut tensor, mapper) = param.consume();
        let path = self.path.join(".");
        if self.group.matches(&id, Some(&path)) {
            let range = compute_range(&self.scheme, &tensor, &self.calibration);
            let qparams = compute_q_params(&self.scheme, range);
            tensor = tensor.quantize(&self.scheme, qparams);
        }
        Param::from_mapped_value(id, tensor, mapper)
    }
}

#[cfg(all(test, not(feature = "tch")))]
mod tests {
    use crate::module::{Module, ParamGroup, Quantizer};
    use crate::tensor::DType;
    use crate::test_device;
    use crate::test_utils::SimpleLinear;
    use burn_tensor::{
        Device, Tensor, Tolerance,
        quantization::{Calibration, QuantLevel, QuantParam, QuantScheme, QuantValue},
    };

    /// Per-tensor Q8 symmetric scheme used across these tests.
    fn test_scheme(device: &Device) -> QuantScheme {
        device
            .settings()
            .quantization
            .scheme
            .with_value(QuantValue::Q8S)
            .with_level(QuantLevel::Tensor)
            .with_param(QuantParam::F32)
    }

    /// Whether a tensor currently holds quantized (`QFloat`) data.
    fn is_quantized<const D: usize>(tensor: &Tensor<D>) -> bool {
        matches!(tensor.dtype(), DType::QFloat(_))
    }

    #[test]
    fn should_quantize_module() {
        let device = test_device();
        let module = SimpleLinear::new(32, 32, &device);
        let scheme = test_scheme(&device);

        let result = module.weight.val();

        let calibration = Calibration::MinMax;
        let mut quantizer = Quantizer::new(calibration, scheme);
        let q_module = module.quantize_weights(&mut quantizer);
        let q_result = q_module.weight.val().dequantize();

        result
            .into_data()
            .assert_approx_eq::<f32>(&q_result.into_data(), Tolerance::permissive());
    }

    #[test]
    fn should_quantize_only_group() {
        let device = test_device();
        let module = SimpleLinear::new(32, 32, &device);
        let scheme = test_scheme(&device);

        let weight_ref = module.weight.val();

        let mut quantizer = Quantizer::new(Calibration::MinMax, scheme);
        let q_module =
            module.quantize_weights_group(&mut quantizer, ParamGroup::from_path("weight"));

        // Only the weight (the group) is quantized; the bias is left untouched.
        assert!(is_quantized(&q_module.weight.val()));
        let bias = q_module.bias.clone().expect("bias should be present");
        assert!(!is_quantized(&bias.val()));

        // The quantized weight still approximates the original values.
        weight_ref.into_data().assert_approx_eq::<f32>(
            &q_module.weight.val().dequantize().into_data(),
            Tolerance::permissive(),
        );
    }

    #[test]
    fn should_quantize_all_when_no_group() {
        let device = test_device();
        let module = SimpleLinear::new(32, 32, &device);
        let scheme = test_scheme(&device);

        let mut quantizer = Quantizer::new(Calibration::MinMax, scheme);
        let q_module = module.quantize_weights(&mut quantizer);

        // Without a group every float parameter is quantized.
        assert!(is_quantized(&q_module.weight.val()));
        assert!(is_quantized(&q_module.bias.clone().unwrap().val()));
    }
}