better val help/fix model improve check

Author: mluerig

bioencoder/scripts/split_dataset.py

@@ -33,10 +33,10 @@ def split_dataset(
     image_dir : str
         Path to the directory containing subfolders of images, where each subfolder represents a class.
     mode : str, optional
-        Specifies the strategy for splitting the dataset:
-            - "flat": Calculating split to the most abundant class (after applying max_ratio), and then applying it to all classes 
-            - "random": Randomly selects images across all classes to form the validation set, disregarding class balance.
-            - "fixed": Ensures each class contributes a fixed proportion to the validation set, based on `val_percent`.
+        Strategy for populating the validation subset:
+            - "flat": Derives a single validation quota from the capped largest class (max_ratio), and applies it uniformly to all classes.
+            - "random": Builds the validation set by drawing images uniformly at random from the pooled, balanced dataset, ignoring class membership.
+            - "fixed": Assigns each class its own `val_percent` share to the validation set based on its balanced size.
         Default is "flat".
     val_percent : float, optional
         Proportion of the dataset to allocate to the validation set, expressed as a decimal.

bioencoder/scripts/train.py

@@ -75,6 +75,7 @@ def train(
     ema_decay_per_epoch = hyperparams["train"]["ema_decay_per_epoch"]
     n_epochs = hyperparams["train"]["n_epochs"]
     target_metric = hyperparams["train"]["target_metric"]
+    min_improvement = hyperparams["train"].get("min_improvement", 0.01)
     stage = hyperparams["train"]["stage"]
     optimizer_params = hyperparams["optimizer"]
     scheduler_params = hyperparams["scheduler"]
@@ -319,15 +320,14 @@ def train(
                     pass
 
             # check if the best value of metric changed. If so -> save the model
-            if (
-                valid_metrics[target_metric] > metric_best*0.99
-            ):  # > 0 if wanting to save all models 
+            current_metric = valid_metrics[target_metric]
+            if metric_best == 0 or current_metric > metric_best * (1 + min_improvement):
                 logger.info(
-                    "{} increased ({:.6f} --> {:.6f}).  Saving model ...".format(
-                        target_metric, metric_best, valid_metrics[target_metric]
+                    "{} improved by ≥{:.2%} ({:.6f} --> {:.6f}). Saving model ...".format(
+                        target_metric, min_improvement, metric_best, current_metric
                     )
                 )
-    
+   
                 torch.save(
                     {
                         "epoch": epoch,
@@ -336,8 +336,10 @@ def train(
                     },
                     os.path.join(weights_dir, f"epoch{epoch}"),
                 )
-                metric_best = valid_metrics[target_metric]
-    
+                metric_best = current_metric
+            else:
+                logger.info(f"Metric {target_metric} did not improve by ≥{min_improvement:.2%} (best: {metric_best:.6f}, current: {current_metric:.6f})")
+
             # if ema is used, go back to regular weights without ema
             if ema:
                 utils.copy_parameters_to_model(copy_of_model_parameters, model)