burn_optim/grad_clipping/
base.rs

1use burn_core as burn;
2
3use burn::{config::Config, tensor::Tensor};
4
5/// Gradient Clipping provides a way to mitigate exploding gradients
6#[derive(Config, Debug)]
7pub enum GradientClippingConfig {
8    /// Clip the gradient by value.
9    Value(f32),
10
11    /// Clip the gradient by norm.
12    Norm(f32),
13}
14
15impl GradientClippingConfig {
16    /// Initialize the gradient clipping.
17    ///
18    /// # Returns
19    ///
20    /// The gradient clipping.
21    pub fn init(&self) -> GradientClipping {
22        match self {
23            GradientClippingConfig::Value(val) => GradientClipping::Value(*val),
24            GradientClippingConfig::Norm(val) => GradientClipping::Norm(*val),
25        }
26    }
27}
28
29/// Gradient Clipping provides a way to mitigate exploding gradients
30/// by clipping every component of the gradient by value or by norm during
31/// backpropagation.
32#[derive(Clone)]
33pub enum GradientClipping {
34    /// Clip the gradient by value.
35    Value(f32),
36
37    /// Clip the gradient by norm.
38    Norm(f32),
39}
40
41impl GradientClipping {
42    /// Clip the gradient.
43    ///
44    /// # Arguments
45    ///
46    /// * `grad` - The gradient to clip.
47    ///
48    /// # Returns
49    ///
50    /// The clipped gradient.
51    pub fn clip_gradient<const D: usize>(&self, grad: Tensor<D>) -> Tensor<D> {
52        match self {
53            GradientClipping::Value(threshold) => self.clip_by_value(grad, *threshold),
54            GradientClipping::Norm(max_norm) => self.clip_by_norm(grad, *max_norm),
55        }
56    }
57
58    fn clip_by_value<const D: usize>(&self, grad: Tensor<D>, threshold: f32) -> Tensor<D> {
59        let greater_mask = grad.clone().greater_scalar(threshold);
60        let lower_mask = grad.clone().lower_scalar(-threshold);
61
62        let clipped_grad = grad.mask_fill(greater_mask, threshold);
63
64        clipped_grad.mask_fill(lower_mask, -threshold)
65    }
66
67    fn clip_by_norm<const D: usize>(&self, grad: Tensor<D>, threshold: f32) -> Tensor<D> {
68        let norm = Self::l2_norm(grad.clone());
69        let min_positive = grad
70            .dtype()
71            .finfo()
72            .unwrap_or(burn::tensor::FloatDType::F32.finfo())
73            .min_positive;
74        let clip_coef = threshold / norm.add_scalar(min_positive);
75        let clip_coef_clamped = clip_coef.clamp_max(1.0);
76        grad.mul(clip_coef_clamped.unsqueeze())
77    }
78
79    fn l2_norm<const D: usize>(tensor: Tensor<D>) -> Tensor<1> {
80        let squared = tensor.square();
81        let sum = squared.sum();
82        sum.sqrt()
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use burn::tensor::Tensor;
90
91    #[test]
92    fn test_clip_by_value() {
93        let gradient: Tensor<2> = Tensor::from_floats(
94            [
95                [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310],
96                [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883],
97            ],
98            &Default::default(),
99        );
100
101        let clipped_gradient = GradientClipping::Value(0.5).clip_gradient(gradient);
102        let clipped_gradient_data = clipped_gradient.into_data();
103
104        for value in clipped_gradient_data.iter::<f32>() {
105            assert!(value <= 0.5);
106        }
107    }
108
109    #[test]
110    fn test_clip_by_norm() {
111        let gradient: Tensor<2> = Tensor::from_floats(
112            [
113                [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310],
114                [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883],
115            ],
116            &Default::default(),
117        );
118
119        let clipped_gradient = GradientClipping::Norm(2.2).clip_gradient(gradient);
120        let clipped_gradient_data = clipped_gradient.into_data();
121
122        for value in clipped_gradient_data.iter::<f32>() {
123            assert!(value <= 0.88);
124        }
125    }
126    #[test]
127    fn test_clip_by_norm_no_clipping() {
128        let gradient: Tensor<2> = Tensor::from_floats(
129            [[0.3, 0.4, 0.5, 0.2], [0.1, 0.6, 0.3, 0.4]],
130            &Default::default(),
131        );
132
133        let clipped_gradient = GradientClipping::Norm(2.2).clip_gradient(gradient.clone());
134
135        clipped_gradient
136            .into_data()
137            .assert_eq(&gradient.into_data(), true);
138    }
139}
burn_optim/grad_clipping/base.rs

burn_optim/grad_clipping/
base.rs
