Add SGD Tests

tracel-ai · nathanielsimard · Sep 18, 2022 · Sep 17, 2022 · Sep 17, 2022 · Sep 17, 2022
commit 192b5dd8113b2a171c321793844ca7ea964f00e1
diff --git a/burn/src/module/state.rs b/burn/src/module/state.rs
@@ -142,6 +142,10 @@ impl<E: Element> StateNamed<E> {
         self.values.get(name)
     }
 
+    pub fn is_empty(&self) -> bool {
+        self.values.is_empty()
+    }
+
     pub fn convert<O: Element>(self) -> StateNamed<O> {
         let mut values = HashMap::with_capacity(self.values.len());
 
@@ -161,6 +165,14 @@ impl<E: Element> State<E> {
         }
     }
 
+    pub fn is_empty(&self) -> bool {
+        match self {
+            State::StateNamed(named) => named.is_empty(),
+            State::Data(_) => false,
+            State::ParamId(_) => false,
+        }
+    }
+
     pub fn convert<O: Element>(self) -> State<O> {
         match self {
             State::StateNamed(named) => State::StateNamed(named.convert()),

diff --git a/burn/src/optim/sgd.rs b/burn/src/optim/sgd.rs
@@ -92,3 +92,100 @@ impl<B: ADBackend> Optimizer for Sgd<B> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        module::{Forward, Module},
+        nn::{Linear, LinearConfig},
+        tensor::{Distribution, Shape},
+        TestADBackend,
+    };
+
+    #[test]
+    fn with_updated_params_should_have_state() {
+        let mut layer = layer();
+        let mut optim = sgd_with_all();
+        let loss = layer.forward(random_tensor());
+        let grads = loss.backward();
+        layer.update_params(&grads, &mut optim);
+
+        let state = optim.state(&layer);
+
+        assert!(!state.is_empty());
+    }
+
+    #[test]
+    fn without_updated_params_should_not_have_state() {
+        let layer = layer();
+        let optim = sgd_with_all();
+
+        let state = optim.state(&layer);
+
+        assert!(state.is_empty());
+    }
+
+    #[test]
+    fn without_momentum_and_weights_decay_should_not_have_state() {
+        let mut layer = layer();
+        let mut optim = sgd_with_nothing();
+        let loss = layer.forward(random_tensor());
+        let grads = loss.backward();
+        layer.update_params(&grads, &mut optim);
+
+        let state = optim.state(&layer);
+
+        assert!(state.is_empty());
+    }
+
+    #[test]
+    fn should_load_state() {
+        let mut layer = layer();
+        let mut optim = sgd_with_all();
+        let loss = layer.forward(random_tensor());
+        let grads = loss.backward();
+        layer.update_params(&grads, &mut optim);
+
+        let state = optim.state(&layer);
+        let mut optim_new = sgd_with_all();
+        let state_new = optim_new.state(&layer);
+        optim_new.load(&layer, &state).unwrap();
+        let state_restored = optim_new.state(&layer);
+
+        assert_ne!(state, state_new);
+        assert_eq!(state, state_restored);
+    }
+
+    fn random_tensor() -> Tensor<TestADBackend, 2> {
+        Tensor::<TestADBackend, 2>::random(Shape::new([2, 20]), Distribution::Standard)
+    }
+
+    fn layer() -> Linear<TestADBackend> {
+        Linear::<TestADBackend>::new(&LinearConfig {
+            d_input: 20,
+            d_output: 20,
+            bias: true,
+        })
+    }
+
+    fn sgd_with_all() -> Sgd<TestADBackend> {
+        Sgd::new(&SgdConfig {
+            learning_rate: 0.02,
+            weight_decay: Some(WeightDecayConfig { penalty: 0.05 }),
+            momentum: Some(MomentumConfig {
+                momentum: 0.9,
+                dampening: 0.1,
+                nesterov: true,
+            }),
+        })
+    }
+
+    fn sgd_with_nothing() -> Sgd<TestADBackend> {
+        Sgd::new(&SgdConfig {
+            learning_rate: 0.02,
+            weight_decay: None,
+            momentum: None,
+        })
+    }
+}