[fix] early stop error when specify metric in training (#333)

* [fix] early stop error when specify metric in training

* [fix] enable find_unused_parameters=True for ddp to avoid exception in unimolv2

* [typo] fix typo

Author: emotionor

unimol_tools/setup.py

@@ -5,8 +5,8 @@
 
 setup(
     name="unimol_tools",
-    version="0.1.3",
-    description=("unimol_tools is a Python package for property prediciton with Uni-Mol in molecule, materials and protein."),
+    version="0.1.3.post1",
+    description=("unimol_tools is a Python package for property prediction with Uni-Mol in molecule, materials and protein."),
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',
     author="DP Technology",

unimol_tools/unimol_tools/tasks/trainer.py

@@ -205,7 +205,7 @@ def fit_predict(
                     # print(f"Main function returned: {y_preds}")
                 except:
                     print("No return value received from main function.")
-                return y_preds
+            return y_preds
         else:
             return self.fit_predict_wo_ddp(
                 model,
@@ -341,7 +341,7 @@ def fit_predict_with_ddp(
         """
         self.init_ddp(local_rank)
         model = model.to(local_rank)
-        model = DistributedDataParallel(model, device_ids=[local_rank])
+        model = DistributedDataParallel(model, device_ids=[local_rank], find_unused_parameters=True)
         train_dataloader = NNDataLoader(
             feature_name=feature_name,
             dataset=train_dataset,
@@ -719,7 +719,7 @@ def inference_with_ddp(
         """
         self.init_ddp(local_rank)
         model = model.to(local_rank)
-        model = DistributedDataParallel(model, device_ids=[local_rank])
+        model = DistributedDataParallel(model, device_ids=[local_rank], find_unused_parameters=True)
         dataloader = NNDataLoader(
             feature_name=feature_name,
             dataset=dataset,
@@ -870,6 +870,7 @@ def __init__(self, patience, dump_dir, fold, metrics, metrics_str):
         self.metrics_str = metrics_str
         self.wait = 0
         self.min_loss = float("inf")
+        self.max_loss = float("-inf")
         self.is_early_stop = False
 
     def early_stop_choice(self, model, epoch, loss, metric_score=None):
@@ -890,16 +891,22 @@ def early_stop_choice(self, model, epoch, loss, metric_score=None):
         ]:
             return self._judge_early_stop_loss(loss, model, epoch)
         else:
-            return self.metrics._early_stop_choice(
+            is_early_stop, min_score, wait, max_score = self.metrics._early_stop_choice(
                 self.wait,
                 self.min_loss,
                 metric_score,
+                self.max_loss,
                 model,
                 self.dump_dir,
                 self.fold,
                 self.patience,
                 epoch,
             )
+            self.min_loss = min_score
+            self.max_loss = max_score
+            self.wait = wait
+            self.is_early_stop = is_early_stop
+            return self.is_early_stop
 
     def _judge_early_stop_loss(self, loss, model, epoch):
         """