Feat/RMSProp-optimizer (#607)

Author: AuruTus

burn-core/src/module/base.rs

@@ -176,8 +176,8 @@ pub trait Module<B: Backend>: Clone + Send + Sync + core::fmt::Debug {
 
     /// Map each tensor in the module with a [mapper](ModuleMapper).
     fn map<M: ModuleMapper<B>>(self, mapper: &mut M) -> Self;
-    /// Load the module state from a record.
 
+    /// Load the module state from a record.
     fn load_record(self, record: Self::Record) -> Self;
 
     /// Convert the module into a record containing the state.

burn-core/src/optim/decay.rs

@@ -16,7 +16,7 @@ pub struct WeightDecayConfig {
 /// State of [WeightDecay](WeightDecay).
 #[derive(Record, Clone, new)]
 pub struct WeightDecayState<B: Backend, const D: usize> {
-    grad_last_step: Tensor<B, D>,
+    pub(crate) grad_last_step: Tensor<B, D>,
 }
 
 /// Weight decay implementation that transforms gradients.
@@ -57,6 +57,15 @@ impl<B: Backend> WeightDecay<B> {
 
         (grad, WeightDecayState::new(grad_last_step))
     }
+
+    /// temp fix for Transform.
+    pub fn transform_temp_fix<const D: usize>(
+        &self,
+        grad: Tensor<B, D>,
+        tensor: Tensor<B, D>,
+    ) -> Tensor<B, D> {
+        tensor.mul_scalar(self.penalty).add(grad)
+    }
 }
 
 impl<B: Backend, const D: usize> WeightDecayState<B, D> {

burn-core/src/optim/mod.rs

@@ -10,6 +10,7 @@ mod adamw;
 mod base;
 mod grad_accum;
 mod grads;
+mod rmsprop;
 mod sgd;
 mod simple;
 mod visitor;
@@ -20,5 +21,6 @@ pub use adamw::*;
 pub use base::*;
 pub use grad_accum::*;
 pub use grads::*;
+pub use rmsprop::*;
 pub use sgd::*;
 pub use simple::*;