burn_optim/optim/
sgd.rs

1use burn_core as burn;
2
3use super::Optimizer;
4use super::decay::{WeightDecay, WeightDecayConfig};
5use super::module_optimizer::ModuleOptimizer;
6use super::momentum::{Momentum, MomentumConfig, MomentumState};
7use crate::LearningRate;
8use crate::RecordState;
9use crate::grad_clipping::GradientClippingConfig;
10use burn::config::Config;
11use burn::tensor::Device;
12use burn::tensor::Tensor;
13
14/// Configuration to create the [Sgd](Sgd) optimizer.
15#[derive(Config, Debug)]
16pub struct SgdConfig {
17    /// [Weight decay](WeightDecayConfig) config.
18    weight_decay: Option<WeightDecayConfig>,
19    /// [Momentum](MomentumConfig) config.
20    momentum: Option<MomentumConfig>,
21    /// [Gradient Clipping](GradientClippingConfig) config.
22    gradient_clipping: Option<GradientClippingConfig>,
23}
24
25/// Optimizer that implements stochastic gradient descent with momentum.
26///
27/// The optimizer can be configured with [SgdConfig](SgdConfig).
28#[derive(Clone)]
29pub struct Sgd {
30    momentum: Option<Momentum>,
31    weight_decay: Option<WeightDecay>,
32}
33
34/// State of [Sgd](Sgd).
35#[derive(RecordState, Clone, new)]
36pub struct SgdState<const D: usize> {
37    /// The current state of the momentum (if any).
38    pub momentum: Option<MomentumState<D>>,
39}
40
41impl SgdConfig {
42    /// Build a [`Sgd`] from the config.
43    ///
44    /// The bare optimizer, which
45    /// [`ModuleOptimizer::with_group`](crate::ModuleOptimizer::with_group) takes to
46    /// optimize one parameter group. [`init`](Self::init) is the whole-module
47    /// counterpart, and the only one that applies the configured gradient clipping.
48    pub fn build(&self) -> Sgd {
49        Sgd {
50            momentum: self.momentum.as_ref().map(Momentum::new),
51            weight_decay: self.weight_decay.as_ref().map(WeightDecay::new),
52        }
53    }
54
55    /// Initializes the SGD optimizer from the configuration.
56    pub fn init(&self) -> ModuleOptimizer {
57        let mut optim = ModuleOptimizer::from(self.build());
58        if let Some(config) = &self.gradient_clipping {
59            optim = optim.with_grad_clipping(config.init());
60        }
61        optim
62    }
63}
64
65impl Optimizer for Sgd {
66    type State<const D: usize> = SgdState<D>;
67
68    fn step<const D: usize>(
69        &self,
70        lr: LearningRate,
71        tensor: Tensor<D>,
72        mut grad: Tensor<D>,
73        state: Option<Self::State<D>>,
74    ) -> (Tensor<D>, Option<Self::State<D>>) {
75        let mut state_momentum = None;
76
77        if let Some(state) = state {
78            state_momentum = state.momentum;
79        }
80
81        if let Some(weight_decay) = &self.weight_decay {
82            grad = weight_decay.transform(grad, tensor.clone());
83        }
84
85        if let Some(momentum) = &self.momentum {
86            let (grad_out, state) = momentum.transform(grad, state_momentum);
87            state_momentum = Some(state);
88            grad = grad_out;
89        }
90
91        let state = SgdState::new(state_momentum);
92        let delta = grad.mul_scalar(lr);
93
94        (tensor - delta, Some(state))
95    }
96
97    fn to_device<const D: usize>(mut state: Self::State<D>, device: &Device) -> Self::State<D> {
98        state.momentum = state.momentum.map(|state| state.to_device(device));
99        state
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use crate::optim::test_utils::assert_optimizer_resume;
107    use crate::{grad_clipping::GradientClipping, optim::GradientsParams};
108    use burn::tensor::{Distribution, Shape};
109    use burn_nn::{Linear, LinearConfig};
110
111    const LEARNING_RATE: LearningRate = 0.02;
112
113    #[test]
114    fn with_updated_params_should_have_state() {
115        let device = Device::default().autodiff();
116        let layer = layer(&device);
117        let mut optim = sgd_with_all();
118        let loss = layer.forward(random_tensor(&device));
119        let grads = loss.backward();
120        let grads = GradientsParams::from_grads(grads, &layer);
121        let _layer = optim.step(LEARNING_RATE, layer, grads);
122
123        let record = optim.to_record();
124
125        assert!(!record.is_empty());
126    }
127
128    #[test]
129    fn without_updated_params_should_not_have_state() {
130        let optim = sgd_with_all();
131        let record = optim.to_record();
132        assert!(record.is_empty());
133    }
134
135    #[test]
136    fn can_attach_gradient_clipping() {
137        let optim = sgd_with_all().with_grad_clipping(GradientClipping::Value(0.5));
138        assert!(optim.has_gradient_clipping());
139    }
140
141    #[test]
142    fn should_load_state() {
143        let device = Device::default().autodiff();
144        assert_optimizer_resume(sgd_with_all, layer(&device), LEARNING_RATE);
145    }
146
147    #[test]
148    fn lora_finetune_trains_adapter_and_freezes_base() {
149        use burn::module::{Lora, Module};
150        use burn::tensor::Tolerance;
151
152        let device = Device::default().autodiff();
153        let linear = LinearConfig::new(8, 8).init(&device);
154
155        // Snapshot the base weight before applying LoRA / training.
156        let base_before = linear.weight.val();
157
158        // Apply LoRA transparently — no change to `Linear` or its forward code.
159        let mut model = linear.apply_lora(Lora::new(4, 8.0));
160        assert!(model.weight.adapter().is_some());
161        let b_before = model.weight.adapter().unwrap().b.val();
162
163        let x = Tensor::<2>::random(Shape::new([16, 8]), Distribution::Default, &device);
164        let target = Tensor::<2>::random(Shape::new([16, 8]), Distribution::Default, &device);
165
166        let mut optim = SgdConfig::new().init();
167
168        let mut first_loss = None;
169        let mut last_loss = 0.0f32;
170        for _ in 0..30 {
171            let output = model.forward(x.clone());
172            let loss = (output - target.clone()).powf_scalar(2.0).mean();
173            last_loss = loss.clone().into_scalar::<f32>();
174            first_loss.get_or_insert(last_loss);
175
176            let grads = loss.backward();
177            let grads = GradientsParams::from_grads(grads, &model);
178            model = optim.step(0.5, model, grads);
179        }
180
181        // Training reduced the loss...
182        assert!(
183            last_loss < first_loss.unwrap(),
184            "expected loss to decrease ({} -> {})",
185            first_loss.unwrap(),
186            last_loss
187        );
188
189        // ...the adapter is still attached and its B factor moved away from the zero init...
190        assert!(model.weight.adapter().is_some());
191        let b_after = model.weight.adapter().unwrap().b.val();
192        let b_change = (b_after - b_before).abs().sum().into_scalar::<f32>();
193        assert!(
194            b_change > 0.0,
195            "adapter factor B should be updated by training"
196        );
197
198        // ...while the frozen base weight is left untouched.
199        model
200            .weight
201            .base()
202            .into_data()
203            .assert_approx_eq::<f32>(&base_before.into_data(), Tolerance::default());
204    }
205
206    fn random_tensor(device: &Device) -> Tensor<2> {
207        Tensor::<2>::random(Shape::new([2, 20]), Distribution::Default, device)
208    }
209
210    fn layer(device: &Device) -> Linear {
211        LinearConfig::new(20, 20).init(device)
212    }
213
214    fn sgd_with_all() -> ModuleOptimizer {
215        SgdConfig {
216            weight_decay: Some(WeightDecayConfig { penalty: 0.05 }),
217            momentum: Some(MomentumConfig {
218                momentum: 0.9,
219                dampening: 0.1,
220                nesterov: true,
221            }),
222            gradient_clipping: None,
223        }
224        .init()
225    }
226}
burn_optim/optim/sgd.rs

burn_optim/optim/
sgd.rs
