Feat/log loss for pretrain (#46)

* Feat/logloss for pretrain * fine-tune L1 weight * update version
open-spaced-repetition · Sep 26, 2023 · 94be1a2 · 94be1a2
1 parent 68df564
commit 94be1a2
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "FSRS-Optimizer"
-version = "4.14.3"
+version = "4.15.0"
 readme = "README.md"
 dependencies = [
     "matplotlib>=3.7.0",

diff --git a/src/fsrs_optimizer/fsrs_optimizer.py b/src/fsrs_optimizer/fsrs_optimizer.py
@@ -854,9 +854,13 @@ def pretrain(self, dataset=None, verbose=True):
 
             def loss(stability):
                 y_pred = power_forgetting_curve(delta_t, stability)
-                rmse = np.sqrt(np.sum((recall - y_pred) ** 2 * count) / total_count)
-                l1 = np.abs(stability - init_s0) / np.sqrt(s0_size) / total_count
-                return rmse + l1
+                logloss = sum(
+                    -(recall * np.log(y_pred) + (1 - recall) * np.log(1 - y_pred))
+                    * count
+                    / total_count
+                )
+                l1 = np.abs(stability - init_s0) / total_count / 16
+                return logloss + l1
 
             res = minimize(
                 loss,