Feat/early stopping feature implementation

etna_core/src/lib.rs

@@ -48,7 +48,7 @@ impl EtnaModel {
         }
     }
 
-    #[pyo3(signature = (x, y, epochs, lr, batch_size=32, weight_decay=0.0, optimizer="sgd", progress_callback=None))]
+    #[pyo3(signature = (x, y, epochs, lr, batch_size=32, weight_decay=0.0, optimizer="sgd", early_stopping=false, patience=10, restore_best=true, progress_callback=None))]
     #[allow(clippy::too_many_arguments)]
     fn train(
         &mut self,
@@ -60,6 +60,9 @@ impl EtnaModel {
         batch_size: usize,
         weight_decay: f32,
         optimizer: &str,
+        early_stopping: bool,
+        patience: usize,
+        restore_best: bool,
         progress_callback: Option<Py<PyAny>>,
     ) -> PyResult<Vec<f32>> {
         let x_vec = pylist_to_vec2(x)?;
@@ -86,6 +89,9 @@ impl EtnaModel {
             weight_decay,
             optimizer_type,
             batch_size,
+            early_stopping,
+            patience,
+            restore_best,
             callback,
         );
 

etna_core/src/model.rs

@@ -108,6 +108,9 @@ impl SimpleNN {
             weight_decay,
             optimizer_type,
             batch_size,
+            false, // early_stopping
+            10,    // patience (unused when early_stopping=false)
+            true,  // restore_best (unused when early_stopping=false)
             |epoch, total, loss| {
                 if epoch % 10 == 0 {
                     println!("Epoch {}/{} - Loss: {:.4}", epoch, total, loss);
@@ -127,12 +130,20 @@ impl SimpleNN {
         weight_decay: f32,
         optimizer_type: OptimizerType,
         batch_size: usize,
+        early_stopping: bool,
+        patience: usize,
+        restore_best: bool,
         progress_callback: F,
     ) -> Vec<f32>
     where
         F: Fn(usize, usize, f32),
     {
         let mut loss_history = Vec::new();
+        let patience = if early_stopping && patience == 0 { 1 } else { patience };
+        let mut best_loss = f32::INFINITY;
+        let mut best_epoch = 0usize;
+        let mut bad_epochs = 0usize;
+        let mut best_state_json: Option<String> = None;
 
         // Initialize optimizers ONLY if they don't exist yet
         if self.optimizers.is_empty() {
@@ -210,6 +221,38 @@ impl SimpleNN {
 
             // Call the progress callback instead of printing
             progress_callback(epoch, epochs, avg_loss);
+
+            if early_stopping {
+                // Treat a strictly lower loss as an improvement (with tiny epsilon for float noise)
+                if avg_loss < best_loss - 1e-12 {
+                    best_loss = avg_loss;
+                    best_epoch = epoch;
+                    bad_epochs = 0;
+
+                    if restore_best {
+                        // Snapshot full model state (layers + optimizer state)
+                        best_state_json = serde_json::to_string(self).ok();
+                    }
+                } else {
+                    bad_epochs += 1;
+
+                    if bad_epochs >= patience {
+                        println!(
+                            "Early stopping triggered at epoch {} (best epoch {} with loss {:.6}).",
+                            epoch, best_epoch, best_loss
+                        );
+                        break;
+                    }
+                }
+            }
+        }
+
+        if early_stopping && restore_best {
+            if let Some(json) = best_state_json {
+                if let Ok(best_model) = serde_json::from_str::<SimpleNN>(&json) {
+                    *self = best_model;
+                }
+            }
         }
 
         loss_history

etna_core/src/optimizer.rs

@@ -165,4 +165,4 @@ mod tests {
         assert!((bias[0] - expected_bias).abs() < 1e-5,
             "Bias update incorrect. Got {}, expected approx {}", bias[0], expected_bias);
     }
-}
+}