From ec2eae5f9695ab222d893e490835eacf8279253d Mon Sep 17 00:00:00 2001
From: Jarrett Ye <jarrett.ye@outlook.com>
Date: Sun, 10 Mar 2024 19:40:21 +0800
Subject: [PATCH] Fix/progress goes backwards (#166)

Co-authored-by: Asuka Minato <i@asukaminato.eu.org>
---
 Cargo.lock      |  2 +-
 Cargo.toml      |  2 +-
 src/training.rs | 60 ++++++++++++++++++++++++++++++++-----------------
 3 files changed, 42 insertions(+), 22 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7fc3a0b..a63302a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1055,7 +1055,7 @@ dependencies = [
 
 [[package]]
 name = "fsrs"
-version = "0.5.0"
+version = "0.5.1"
 dependencies = [
  "burn",
  "chrono",
diff --git a/Cargo.toml b/Cargo.toml
index 7539c31..bc90171 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "fsrs"
-version = "0.5.0"
+version = "0.5.1"
 authors = ["Open Spaced Repetition"]
 categories = ["algorithms", "science"]
 edition = "2021"
diff --git a/src/training.rs b/src/training.rs
index 394ebc4..4825c55 100644
--- a/src/training.rs
+++ b/src/training.rs
@@ -214,9 +214,6 @@ impl<B: Backend> FSRS<B> {
         };
 
         let n_splits = 5;
-        if let Some(progress) = &progress {
-            progress.lock().unwrap().splits = vec![ProgressState::default(); n_splits];
-        }
         let average_recall = calculate_average_recall(&items);
         let (pre_trainset, trainsets, testset) = split_data(items, n_splits);
         let initial_stability = pretrain(pre_trainset, average_recall).map_err(|e| {
@@ -239,18 +236,32 @@ impl<B: Backend> FSRS<B> {
             AdamConfig::new(),
         );
 
-        let weight_sets: Result<Vec<Vec<f32>>> = (0..n_splits)
+        let trainsets: Vec<Vec<FSRSItem>> = (0..n_splits)
             .into_par_iter()
             .map(|i| {
-                let trainset = trainsets
+                trainsets
                     .par_iter()
                     .enumerate()
                     .filter(|&(j, _)| j != i)
                     .flat_map(|(_, trainset)| trainset.clone())
-                    .collect();
+                    .collect()
+            })
+            .collect();
+
+        if let Some(progress) = &progress {
+            let mut progress_states = vec![ProgressState::default(); n_splits];
+            for (i, progress_state) in progress_states.iter_mut().enumerate() {
+                progress_state.epoch_total = config.num_epochs;
+                progress_state.items_total = trainsets[i].len();
+            }
+            progress.lock().unwrap().splits = progress_states
+        }
 
+        let weight_sets: Result<Vec<Vec<f32>>> = (0..n_splits)
+            .into_par_iter()
+            .map(|i| {
                 let model = train::<Autodiff<B>>(
-                    trainset,
+                    trainsets[i].clone(),
                     testset.clone(),
                     &config,
                     self.device(),
@@ -280,10 +291,8 @@ impl<B: Backend> FSRS<B> {
             .map(|&sum| sum / n_splits as f32)
             .collect();
 
-        for weight in &average_parameters {
-            if !weight.is_finite() {
-                return Err(FSRSError::InvalidInput);
-            }
+        if average_parameters.iter().any(|weight| weight.is_infinite()) {
+            return Err(FSRSError::InvalidInput);
         }
 
         Ok(average_parameters)
@@ -499,25 +508,36 @@ mod tests {
                 thread::sleep(Duration::from_millis(10));
                 let guard = progress.lock().unwrap();
                 finished = guard.finished();
-                info!("progress: {}/{}", guard.current(), guard.total());
+                println!("progress: {}/{}", guard.current(), guard.total());
             }
         });
 
-        if let Some(progress2) = &progress2 {
-            progress2.lock().unwrap().splits = vec![ProgressState::default(); n_splits];
-        }
-
-        let parameters_sets: Vec<Vec<f32>> = (0..n_splits)
+        let trainsets: Vec<Vec<FSRSItem>> = (0..n_splits)
             .into_par_iter()
             .map(|i| {
-                let trainset = trainsets
+                trainsets
                     .par_iter()
                     .enumerate()
                     .filter(|&(j, _)| j != i)
                     .flat_map(|(_, trainset)| trainset.clone())
-                    .collect();
+                    .collect()
+            })
+            .collect();
+
+        if let Some(progress2) = &progress2 {
+            let mut progress_states = vec![ProgressState::default(); n_splits];
+            for (i, progress_state) in progress_states.iter_mut().enumerate() {
+                progress_state.epoch_total = config.num_epochs;
+                progress_state.items_total = trainsets[i].len();
+            }
+            progress2.lock().unwrap().splits = progress_states
+        }
+
+        let parameters_sets: Vec<Vec<f32>> = (0..n_splits)
+            .into_par_iter()
+            .map(|i| {
                 let model = train::<NdArrayAutodiff>(
-                    trainset,
+                    trainsets[i].clone(),
                     items.clone(),
                     &config,
                     device,