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    pub(crate) fn build(&self) -> Sgd {
44        Sgd {
45            momentum: self.momentum.as_ref().map(Momentum::new),
46            weight_decay: self.weight_decay.as_ref().map(WeightDecay::new),
47        }
48    }
49
50    /// Initializes the SGD optimizer from the configuration.
51    pub fn init(&self) -> ModuleOptimizer {
52        let mut optim = ModuleOptimizer::from(self.build());
53        if let Some(config) = &self.gradient_clipping {
54            optim = optim.with_grad_clipping(config.init());
55        }
56        optim
57    }
58}
59
60impl Optimizer for Sgd {
61    type State<const D: usize> = SgdState<D>;
62
63    fn step<const D: usize>(
64        &self,
65        lr: LearningRate,
66        tensor: Tensor<D>,
67        mut grad: Tensor<D>,
68        state: Option<Self::State<D>>,
69    ) -> (Tensor<D>, Option<Self::State<D>>) {
70        let mut state_momentum = None;
71
72        if let Some(state) = state {
73            state_momentum = state.momentum;
74        }
75
76        if let Some(weight_decay) = &self.weight_decay {
77            grad = weight_decay.transform(grad, tensor.clone());
78        }
79
80        if let Some(momentum) = &self.momentum {
81            let (grad_out, state) = momentum.transform(grad, state_momentum);
82            state_momentum = Some(state);
83            grad = grad_out;
84        }
85
86        let state = SgdState::new(state_momentum);
87        let delta = grad.mul_scalar(lr);
88
89        (tensor - delta, Some(state))
90    }
91
92    fn to_device<const D: usize>(mut state: Self::State<D>, device: &Device) -> Self::State<D> {
93        state.momentum = state.momentum.map(|state| state.to_device(device));
94        state
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101    use crate::{grad_clipping::GradientClipping, optim::GradientsParams};
102    use burn::tensor::{Distribution, Shape};
103    use burn_nn::{Linear, LinearConfig};
104
105    const LEARNING_RATE: LearningRate = 0.02;
106
107    #[test]
108    fn with_updated_params_should_have_state() {
109        let device = Device::default().autodiff();
110        let layer = layer(&device);
111        let mut optim = sgd_with_all();
112        let loss = layer.forward(random_tensor(&device));
113        let grads = loss.backward();
114        let grads = GradientsParams::from_grads(grads, &layer);
115        let _layer = optim.step(LEARNING_RATE.into(), layer, grads);
116
117        let record = optim.to_record();
118
119        assert!(!record.is_empty());
120    }
121
122    #[test]
123    fn without_updated_params_should_not_have_state() {
124        let optim = sgd_with_all();
125        let record = optim.to_record();
126        assert!(record.is_empty());
127    }
128
129    #[test]
130    fn can_attach_gradient_clipping() {
131        let optim = sgd_with_all().with_grad_clipping(GradientClipping::Value(0.5));
132        assert!(optim.has_gradient_clipping());
133    }
134
135    #[test]
136    fn should_load_state() {
137        let device = Device::default().autodiff();
138        let layer = layer(&device);
139        let mut optim = sgd_with_all();
140        let loss = layer.forward(random_tensor(&device));
141        let grads = loss.backward();
142        let grads = GradientsParams::from_grads(grads, &layer);
143        let _layer = optim.step(LEARNING_RATE.into(), layer, grads);
144
145        let record = optim.to_record();
146        let bytes = optim.into_bytes().unwrap();
147        let optim_new = sgd_with_all();
148        let record_new = optim_new.to_record();
149        let optim_new = optim_new.from_bytes(bytes).unwrap();
150        let state_restored = optim_new.to_record();
151
152        assert_ne!(record.len(), record_new.len());
153        assert_eq!(record.len(), state_restored.len());
154    }
155
156    #[test]
157    fn lora_finetune_trains_adapter_and_freezes_base() {
158        use burn::module::{LoraConfig, Module};
159        use burn::tensor::Tolerance;
160
161        let device = Device::default().autodiff();
162        let linear = LinearConfig::new(8, 8).init(&device);
163
164        // Snapshot the base weight before applying LoRA / training.
165        let base_before = linear.weight.val();
166
167        // Apply LoRA transparently — no change to `Linear` or its forward code.
168        let mut model = linear.apply_lora(LoraConfig::new(4, 8.0));
169        assert!(model.weight.adapter().is_some());
170        let b_before = model.weight.adapter().unwrap().b.val();
171
172        let x = Tensor::<2>::random(Shape::new([16, 8]), Distribution::Default, &device);
173        let target = Tensor::<2>::random(Shape::new([16, 8]), Distribution::Default, &device);
174
175        let mut optim = SgdConfig::new().init();
176
177        let mut first_loss = None;
178        let mut last_loss = 0.0f32;
179        for _ in 0..30 {
180            let output = model.forward(x.clone());
181            let loss = (output - target.clone()).powf_scalar(2.0).mean();
182            last_loss = loss.clone().into_scalar::<f32>();
183            first_loss.get_or_insert(last_loss);
184
185            let grads = loss.backward();
186            let grads = GradientsParams::from_grads(grads, &model);
187            model = optim.step(0.5.into(), model, grads);
188        }
189
190        // Training reduced the loss...
191        assert!(
192            last_loss < first_loss.unwrap(),
193            "expected loss to decrease ({} -> {})",
194            first_loss.unwrap(),
195            last_loss
196        );
197
198        // ...the adapter is still attached and its B factor moved away from the zero init...
199        assert!(model.weight.adapter().is_some());
200        let b_after = model.weight.adapter().unwrap().b.val();
201        let b_change = (b_after - b_before).abs().sum().into_scalar::<f32>();
202        assert!(
203            b_change > 0.0,
204            "adapter factor B should be updated by training"
205        );
206
207        // ...while the frozen base weight is left untouched.
208        model
209            .weight
210            .base()
211            .into_data()
212            .assert_approx_eq::<f32>(&base_before.into_data(), Tolerance::default());
213    }
214
215    fn random_tensor(device: &Device) -> Tensor<2> {
216        Tensor::<2>::random(Shape::new([2, 20]), Distribution::Default, device)
217    }
218
219    fn layer(device: &Device) -> Linear {
220        LinearConfig::new(20, 20).init(device)
221    }
222
223    fn sgd_with_all() -> ModuleOptimizer {
224        SgdConfig {
225            weight_decay: Some(WeightDecayConfig { penalty: 0.05 }),
226            momentum: Some(MomentumConfig {
227                momentum: 0.9,
228                dampening: 0.1,
229                nesterov: true,
230            }),
231            gradient_clipping: None,
232        }
233        .init()
234    }
235}
burn_optim/optim/sgd.rs

burn_optim/optim/
sgd.rs
