burn_optim/optim/
grad_accum.rs

1use burn_core as burn;
2
3use core::marker::PhantomData;
4
5use burn::module::{AutodiffModule, ModuleVisitor, Param};
6use burn::tensor::Tensor;
7
8use super::GradientsParams;
9
10/// Accumulate gradients into a single [GradientsParams] object.
11pub struct GradientsAccumulator<M> {
12    grads: GradientsParams,
13    phantom: PhantomData<M>,
14}
15
16impl<M> Default for GradientsAccumulator<M> {
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl<M> GradientsAccumulator<M> {
23    /// Create a new gradients accumulator.
24    pub fn new() -> Self {
25        Self {
26            grads: GradientsParams::new(),
27            phantom: PhantomData,
28        }
29    }
30}
31
32impl<M> GradientsAccumulator<M> {
33    /// Accumulate the given gradients for each parameter in the given module.
34    pub fn accumulate(&mut self, module: &M, grads: GradientsParams)
35    where
36        M: AutodiffModule,
37    {
38        let mut visitor = ModuleGradsAccumulator::<M>::new(&mut self.grads, grads);
39        module.visit(&mut visitor);
40    }
41
42    /// Return the accumulated gradients and reset the accumulator state.
43    pub fn grads(&mut self) -> GradientsParams {
44        let mut grads = GradientsParams::new();
45        core::mem::swap(&mut self.grads, &mut grads);
46
47        grads
48    }
49}
50
51#[derive(new)]
52struct ModuleGradsAccumulator<'a, M> {
53    grads: &'a mut GradientsParams,
54    grads_new: GradientsParams,
55    phantom: PhantomData<M>,
56}
57
58impl<M: AutodiffModule> ModuleVisitor for ModuleGradsAccumulator<'_, M> {
59    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<D>>) {
60        let grad_updated = match self.grads_new.remove::<D>(param.id) {
61            Some(new) => match self.grads.remove::<D>(param.id) {
62                Some(grad) => grad.add(new),
63                None => new,
64            },
65            None => match self.grads.remove::<D>(param.id) {
66                Some(grad) => grad,
67                None => return,
68            },
69        };
70
71        self.grads.register::<D>(param.id, grad_updated);
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78    use burn::tensor::{Device, Distribution};
79    use burn_nn::{Linear, LinearConfig};
80
81    #[test]
82    fn test_accumulate_gradients_one_step() {
83        let device = Device::default().autodiff();
84        let mut accumulator = GradientsAccumulator::new();
85        let layer = layer(&device);
86        let loss = layer.forward(random_tensor(&device));
87        let grads = GradientsParams::from_grads(loss.backward(), &layer);
88
89        accumulator.accumulate(&layer, grads);
90
91        let grads = accumulator.grads();
92        assert!(!grads.is_empty())
93    }
94
95    #[test]
96    fn test_accumulate_gradients_two_steps() {
97        let device = Device::default().autodiff();
98        let mut accumulator = GradientsAccumulator::new();
99        let layer = layer(&device);
100        let loss_1 = layer.forward(random_tensor(&device));
101        let loss_2 = layer.forward(random_tensor(&device));
102        let grads_1 = GradientsParams::from_grads(loss_1.backward(), &layer);
103        let grads_2 = GradientsParams::from_grads(loss_2.backward(), &layer);
104
105        accumulator.accumulate(&layer, grads_1);
106        accumulator.accumulate(&layer, grads_2);
107
108        let grads = accumulator.grads();
109        assert_eq!(grads.len(), 2)
110    }
111
112    fn layer(device: &Device) -> Linear {
113        LinearConfig::new(20, 20).init(device)
114    }
115
116    fn random_tensor(device: &Device) -> Tensor<2> {
117        Tensor::<2>::random([2, 20], Distribution::Default, device)
118    }
119}
burn_optim/optim/grad_accum.rs

burn_optim/optim/
grad_accum.rs
