implementing early stop for benchmark and making methods consistent

src/benchmark/methods/abco.py

@@ -48,6 +48,7 @@ def train(
         n_epochs: int,
         train_dataset: Dict[str, DataLoader],
         eval_dataset: Dict[str, DataLoader] = None,
+        always_ave: bool = False,
         folder: str = None
     ) -> Self:
         if folder is None:

src/benchmark/methods/bc.py

@@ -115,18 +115,22 @@ def train(
         train_dataset: DataLoader,
         eval_dataset: DataLoader = None,
         always_save: bool = False,
+        folder: str = None
     ) -> Self:
         """Train process.
 
         Args:
             n_epochs (int): amount of epoch to run.
             train_dataset (DataLoader): data to train.
             eval_dataset (DataLoader): data to eval. Defaults to None.
+            always_save (bool): whether it should save all eval steps.
+            folder (str): a specific folder to save the benchmark results.
 
         Returns:
             method (Self): trained method.
         """
-        folder = f"./benchmark_results/bc/{self.environment_name}"
+        if folder is None:
+            folder = f"./benchmark_results/bc/{self.environment_name}"
         if not os.path.exists(folder):
             os.makedirs(f"{folder}/")
 
@@ -158,7 +162,7 @@ def train(
                     best_model = metrics["aer"]
                     self.save(name=epoch if always_save else None)
 
-                    if early_stop(metrics["aer"]):
+                    if self.early_stop(metrics["aer"]):
                         return self
 
         return self

src/benchmark/methods/bco.py

@@ -146,6 +146,7 @@ def train(
         n_epochs: int,
         train_dataset: Dict[str, DataLoader],
         eval_dataset: Dict[str, DataLoader] = None,
+        always_save: bool = False,
         folder: str = None
     ) -> Self:
         """Train process.
@@ -154,6 +155,8 @@ def train(
             n_epochs (int): amount of epoch to run.
             train_dataset (DataLoader): data to train.
             eval_dataset (DataLoader): data to eval. Defaults to None.
+            always_save (bool): whether it should save all eval steps.
+            folder (str): a specific folder to save the benchmark results.
 
         Returns:
             method (Self): trained method.
@@ -209,12 +212,15 @@ def train(
             if epoch % self.enjoy_criteria == 0:
                 train_dataset = self._append_samples(train_dataset)
 
-            if epoch % self.enjoy_criteria == 0 or epoch + 1 == n_epochs:
+            if epoch > 0 and epoch % self.enjoy_criteria == 0:
                 metrics = self._enjoy()
                 board.add_scalars("Enjoy", epoch="enjoy", **metrics)
                 board.step("enjoy")
-                if best_model < metrics["aer"]:
-                    self.save()
+                if best_model < metrics["aer"] or always_save:
+                    self.save(name=epoch if always_save else None)
+
+                    if self.early_stop(metrics["aer"]):
+                        return self
 
         return self
 

src/benchmark/methods/iupe.py

@@ -93,6 +93,7 @@ def train(
         n_epochs: int,
         train_dataset: Dict[str, DataLoader],
         eval_dataset: Dict[str, DataLoader] = None,
+        always_save: bool = False,
         folder: str = None
     ) -> Self:
         if folder is None:

src/benchmark/methods/method.py

@@ -256,4 +256,4 @@ def early_stop(self, metric: Metrics) -> bool:
         Returns:
             stop (bool): if it should stop or not.
         """
-        raise NotImplementedError()
+        raise False