fix(ml): change early stopping implementation

Early stopping now terminates one epoch later
borkdominik · Oct 15, 2024 · b01d699 · b01d699
1 parent 6b28267
commit b01d699
Show file tree

Hide file tree

Showing 4 changed files with 3 additions and 7 deletions.
diff --git a/ml/gnn/src/gnn.py b/ml/gnn/src/gnn.py
@@ -33,7 +33,7 @@
         "Please provide the train, validation, and test dataset file paths as arguments"
     )
 
-num_epochs = 2000
+num_epochs = 100
 start_epoch = 0
 patience = 10
 

diff --git a/ml/gnn/src/model/base_model.py b/ml/gnn/src/model/base_model.py
@@ -130,7 +130,7 @@ def fit(
                     remaining_patience = patience
                 else:
                     remaining_patience -= 1
-                    if remaining_patience == 0:
+                    if remaining_patience < 0:
                         self.layout_proxy.print(
                             f"{text_padding}Early stopping in epoch {epoch}"
                         )

diff --git a/ml/tree-lstm/src/paper/mdeoperation.py b/ml/tree-lstm/src/paper/mdeoperation.py
@@ -298,7 +298,7 @@ def train(
                 remaining_patience = args.patience
             else:
                 remaining_patience -= 1
-                if remaining_patience == 0:
+                if remaining_patience < 0:
                     print(f"Early stopping in epoch {epoch}")
                     break
             sys.stdout.flush()

diff --git a/ml/tree-lstm/src/tree_dataset.py b/ml/tree-lstm/src/tree_dataset.py
@@ -24,8 +24,6 @@ def load(self):
         dataset_load_start_time = time.perf_counter()
         if self.is_cached:
             self.data, self.metadata, self.vocabulary = torch.load(self.dataset_cache_file)
-            # TODO?
-            # self.to(device)
         else:
             with open(self.dataset_path, "r") as file:
                 dataset_input = json.load(file)
@@ -42,8 +40,6 @@ def load(self):
                     (self.data, self.metadata, self.vocabulary),
                     self.dataset_cache_file,
                 )
-                # TODO?
-                # self.to(device)
 
         dataset_load_end_time = time.perf_counter()
         print(