Feat/option enable_short_term in training (#258)

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "fsrs"
-version = "1.5.0"
+version = "2.0.0"
 authors = ["Open Spaced Repetition"]
 categories = ["algorithms", "science"]
 edition = "2021"

examples/optimize.rs

@@ -18,7 +18,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     println!("Default parameters: {:?}", DEFAULT_PARAMETERS);
 
     // Optimize the FSRS model using the created items
-    let optimized_parameters = fsrs.compute_parameters(fsrs_items, None)?;
+    let optimized_parameters = fsrs.compute_parameters(fsrs_items, None, false)?;
 
     println!("Optimized parameters: {:?}", optimized_parameters);
 

src/model.rs

@@ -41,12 +41,16 @@ impl<B: Backend, const N: usize> Pow<B, N> for Tensor<B, N> {
 impl<B: Backend> Model<B> {
     #[allow(clippy::new_without_default)]
     pub fn new(config: ModelConfig) -> Self {
-        let initial_params = config
+        let mut initial_params: Vec<f32> = config
             .initial_stability
             .unwrap_or_else(|| DEFAULT_PARAMETERS[0..4].try_into().unwrap())
             .into_iter()
             .chain(DEFAULT_PARAMETERS[4..].iter().copied())
             .collect();
+        if config.freeze_short_term_stability {
+            initial_params[17] = 0.0;
+            initial_params[18] = 0.0;
+        }
 
         Self {
             w: Param::from_tensor(Tensor::from_floats(
@@ -199,8 +203,10 @@ pub(crate) struct MemoryStateTensors<B: Backend> {
 #[derive(Config, Module, Debug, Default)]
 pub struct ModelConfig {
     #[config(default = false)]
-    pub freeze_stability: bool,
+    pub freeze_initial_stability: bool,
     pub initial_stability: Option<[f32; 4]>,
+    #[config(default = false)]
+    pub freeze_short_term_stability: bool,
 }
 
 impl ModelConfig {

src/training.rs

@@ -83,6 +83,16 @@ impl<B: AutodiffBackend> Model<B> {
         self.w.grad_replace(&mut grad, updated_grad_tensor);
         grad
     }
+
+    fn free_short_term_stability(&self, mut grad: B::Gradients) -> B::Gradients {
+        let grad_tensor = self.w.grad(&grad).unwrap();
+        let updated_grad_tensor =
+            grad_tensor.slice_assign([17..19], Tensor::zeros([2], &B::Device::default()));
+
+        self.w.grad_remove(&mut grad);
+        self.w.grad_replace(&mut grad, updated_grad_tensor);
+        grad
+    }
 }
 
 #[derive(Debug, Default, Clone)]
@@ -202,6 +212,7 @@ impl<B: Backend> FSRS<B> {
         &self,
         train_set: Vec<FSRSItem>,
         progress: Option<Arc<Mutex<CombinedProgressState>>>,
+        enable_short_term: bool,
     ) -> Result<Vec<f32>> {
         let finish_progress = || {
             if let Some(progress) = &progress {
@@ -235,8 +246,9 @@ impl<B: Backend> FSRS<B> {
         }
         let config = TrainingConfig::new(
             ModelConfig {
-                freeze_stability: false,
+                freeze_initial_stability: !enable_short_term,
                 initial_stability: Some(initial_stability),
+                freeze_short_term_stability: !enable_short_term,
             },
             AdamConfig::new().with_epsilon(1e-8),
         );
@@ -295,7 +307,7 @@ impl<B: Backend> FSRS<B> {
         Ok(optimized_parameters)
     }
 
-    pub fn benchmark(&self, mut train_set: Vec<FSRSItem>) -> Vec<f32> {
+    pub fn benchmark(&self, mut train_set: Vec<FSRSItem>, enable_short_term: bool) -> Vec<f32> {
         let average_recall = calculate_average_recall(&train_set);
         let (pre_train_set, _next_train_set) = train_set
             .clone()
@@ -304,8 +316,9 @@ impl<B: Backend> FSRS<B> {
         let initial_stability = pretrain(pre_train_set, average_recall).unwrap().0;
         let config = TrainingConfig::new(
             ModelConfig {
-                freeze_stability: false,
+                freeze_initial_stability: !enable_short_term,
                 initial_stability: Some(initial_stability),
+                freeze_short_term_stability: !enable_short_term,
             },
             AdamConfig::new().with_epsilon(1e-8),
         );
@@ -377,9 +390,12 @@ fn train<B: AutodiffBackend>(
                 Reduction::Sum,
             );
             let mut gradients = loss.backward();
-            if model.config.freeze_stability {
+            if model.config.freeze_initial_stability {
                 gradients = model.freeze_initial_stability(gradients);
             }
+            if model.config.freeze_short_term_stability {
+                gradients = model.free_short_term_stability(gradients);
+            }
             let grads = GradientsParams::from_grads(gradients, &model);
             model = optim.step(lr, model, grads);
             model.w = parameter_clipper(model.w);
@@ -663,26 +679,30 @@ mod tests {
                 .unwrap();
         }
         for items in [anki21_sample_file_converted_to_fsrs(), data_from_csv()] {
-            let progress = CombinedProgressState::new_shared();
-            let progress2 = Some(progress.clone());
-            thread::spawn(move || {
-                let mut finished = false;
-                while !finished {
-                    thread::sleep(Duration::from_millis(500));
-                    let guard = progress.lock().unwrap();
-                    finished = guard.finished();
-                    println!("progress: {}/{}", guard.current(), guard.total());
-                }
-            });
-
-            let fsrs = FSRS::new(Some(&[])).unwrap();
-            let parameters = fsrs.compute_parameters(items.clone(), progress2).unwrap();
-            dbg!(&parameters);
-
-            // evaluate
-            let model = FSRS::new(Some(&parameters)).unwrap();
-            let metrics = model.evaluate(items, |_| true).unwrap();
-            dbg!(&metrics);
+            for enable_short_term in [true, false] {
+                let progress = CombinedProgressState::new_shared();
+                let progress2 = Some(progress.clone());
+                thread::spawn(move || {
+                    let mut finished = false;
+                    while !finished {
+                        thread::sleep(Duration::from_millis(500));
+                        let guard = progress.lock().unwrap();
+                        finished = guard.finished();
+                        println!("progress: {}/{}", guard.current(), guard.total());
+                    }
+                });
+
+                let fsrs = FSRS::new(Some(&[])).unwrap();
+                let parameters = fsrs
+                    .compute_parameters(items.clone(), progress2, enable_short_term)
+                    .unwrap();
+                dbg!(&parameters);
+
+                // evaluate
+                let model = FSRS::new(Some(&parameters)).unwrap();
+                let metrics = model.evaluate(items.clone(), |_| true).unwrap();
+                dbg!(&metrics);
+            }
         }
     }
 }