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::{grad_clipping::GradientClipping, optim::GradientsParams};
107    use burn::tensor::{Distribution, Shape};
108    use burn_nn::{Linear, LinearConfig};
109
110    const LEARNING_RATE: LearningRate = 0.02;
111
112    #[test]
113    fn with_updated_params_should_have_state() {
114        let device = Device::default().autodiff();
115        let layer = layer(&device);
116        let mut optim = sgd_with_all();
117        let loss = layer.forward(random_tensor(&device));
118        let grads = loss.backward();
119        let grads = GradientsParams::from_grads(grads, &layer);
120        let _layer = optim.step(LEARNING_RATE, layer, grads);
121
122        let record = optim.to_record();
123
124        assert!(!record.is_empty());
125    }
126
127    #[test]
128    fn without_updated_params_should_not_have_state() {
129        let optim = sgd_with_all();
130        let record = optim.to_record();
131        assert!(record.is_empty());
132    }
133
134    #[test]
135    fn can_attach_gradient_clipping() {
136        let optim = sgd_with_all().with_grad_clipping(GradientClipping::Value(0.5));
137        assert!(optim.has_gradient_clipping());
138    }
139
140    #[test]
141    fn should_load_state() {
142        let device = Device::default().autodiff();
143        let layer = layer(&device);
144        let mut optim = sgd_with_all();
145        let loss = layer.forward(random_tensor(&device));
146        let grads = loss.backward();
147        let grads = GradientsParams::from_grads(grads, &layer);
148        let _layer = optim.step(LEARNING_RATE, layer, grads);
149
150        let record = optim.to_record();
151        let bytes = optim.into_bytes().unwrap();
152        let optim_new = sgd_with_all();
153        let record_new = optim_new.to_record();
154        let optim_new = optim_new.from_bytes(bytes).unwrap();
155        let state_restored = optim_new.to_record();
156
157        assert_ne!(record.len(), record_new.len());
158        assert_eq!(record.len(), state_restored.len());
159    }
160
161    #[test]
162    fn lora_finetune_trains_adapter_and_freezes_base() {
163        use burn::module::{Lora, Module};
164        use burn::tensor::Tolerance;
165
166        let device = Device::default().autodiff();
167        let linear = LinearConfig::new(8, 8).init(&device);
168
169        // Snapshot the base weight before applying LoRA / training.
170        let base_before = linear.weight.val();
171
172        // Apply LoRA transparently — no change to `Linear` or its forward code.
173        let mut model = linear.apply_lora(Lora::new(4, 8.0));
174        assert!(model.weight.adapter().is_some());
175        let b_before = model.weight.adapter().unwrap().b.val();
176
177        let x = Tensor::<2>::random(Shape::new([16, 8]), Distribution::Default, &device);
178        let target = Tensor::<2>::random(Shape::new([16, 8]), Distribution::Default, &device);
179
180        let mut optim = SgdConfig::new().init();
181
182        let mut first_loss = None;
183        let mut last_loss = 0.0f32;
184        for _ in 0..30 {
185            let output = model.forward(x.clone());
186            let loss = (output - target.clone()).powf_scalar(2.0).mean();
187            last_loss = loss.clone().into_scalar::<f32>();
188            first_loss.get_or_insert(last_loss);
189
190            let grads = loss.backward();
191            let grads = GradientsParams::from_grads(grads, &model);
192            model = optim.step(0.5, model, grads);
193        }
194
195        // Training reduced the loss...
196        assert!(
197            last_loss < first_loss.unwrap(),
198            "expected loss to decrease ({} -> {})",
199            first_loss.unwrap(),
200            last_loss
201        );
202
203        // ...the adapter is still attached and its B factor moved away from the zero init...
204        assert!(model.weight.adapter().is_some());
205        let b_after = model.weight.adapter().unwrap().b.val();
206        let b_change = (b_after - b_before).abs().sum().into_scalar::<f32>();
207        assert!(
208            b_change > 0.0,
209            "adapter factor B should be updated by training"
210        );
211
212        // ...while the frozen base weight is left untouched.
213        model
214            .weight
215            .base()
216            .into_data()
217            .assert_approx_eq::<f32>(&base_before.into_data(), Tolerance::default());
218    }
219
220    fn random_tensor(device: &Device) -> Tensor<2> {
221        Tensor::<2>::random(Shape::new([2, 20]), Distribution::Default, device)
222    }
223
224    fn layer(device: &Device) -> Linear {
225        LinearConfig::new(20, 20).init(device)
226    }
227
228    fn sgd_with_all() -> ModuleOptimizer {
229        SgdConfig {
230            weight_decay: Some(WeightDecayConfig { penalty: 0.05 }),
231            momentum: Some(MomentumConfig {
232                momentum: 0.9,
233                dampening: 0.1,
234                nesterov: true,
235            }),
236            gradient_clipping: None,
237        }
238        .init()
239    }
240}
burn_optim/optim/sgd.rs

burn_optim/optim/
sgd.rs
