Ensure videos are allocated into all specified splits

tests/conftest.py

@@ -71,6 +71,11 @@ def forward(self, x, *args, **kwargs):
 class DummyTrainConfig(TrainConfig):
     # let model name be "dummy" without causing errors
     model_name: str
+    batch_size = 1
+    max_epochs = 1
+    model_name = "dummy"
+    skip_load_validation = True
+    auto_lr_find = False
 
 
 @pytest.fixture(scope="session")

tests/test_config.py

@@ -203,8 +203,18 @@ def test_labels_with_invalid_split(labels_absolute_path):
 
 
 def test_labels_no_splits(labels_no_splits, tmp_path):
-    config = TrainConfig(data_dir=TEST_VIDEOS_DIR, labels=labels_no_splits, save_dir=tmp_path)
-    assert set(config.labels.split.unique()) == set(("holdout", "train", "val"))
+    labels_three_videos = pd.read_csv(labels_no_splits).head(3)
+    # test with fewer videos and ensure we still get one of each
+    _ = TrainConfig(
+        data_dir=TEST_VIDEOS_DIR,
+        labels=labels_three_videos,
+        save_dir=tmp_path,
+        split_proportions=dict(train=3, val=1, holdout=1),
+    )
+
+    assert set(pd.read_csv(tmp_path / "splits.csv").split.unique()) == set(
+        ["train", "val", "holdout"]
+    )
 
 
 def test_labels_split_proportions(labels_no_splits, tmp_path):

tests/test_model_manager.py

@@ -17,6 +17,20 @@ def test_model_manager(dummy_trainer):
     assert not (dummy_trainer.model.model[3].weight == 0).all()
 
 
+def test_no_early_stopping(
+    labels_absolute_path, tmp_path, dummy_checkpoint, dummy_video_loader_config
+):
+    config = DummyTrainConfig(
+        labels=labels_absolute_path,
+        data_dir=TEST_VIDEOS_DIR,
+        checkpoint=dummy_checkpoint,
+        early_stopping_config=None,
+        save_dir=tmp_path / "my_model",
+        num_workers=1,
+    )
+    train_model(train_config=config, video_loader_config=dummy_video_loader_config)
+
+
 def test_save_checkpoint(dummy_trained_model_checkpoint):
     checkpoint = torch.load(dummy_trained_model_checkpoint)
 
@@ -76,14 +90,9 @@ def test_save_metrics_less_than_two_classes(
         train_config=DummyTrainConfig(
             labels=labels,
             data_dir=TEST_VIDEOS_DIR,
-            model_name="dummy",
             checkpoint=dummy_checkpoint,
-            max_epochs=1,
-            batch_size=1,
-            auto_lr_find=False,
             num_workers=2,
             save_dir=tmp_path / "my_model",
-            skip_load_validation=True,
         ),
         video_loader_config=dummy_video_loader_config,
     )
@@ -136,14 +145,9 @@ def test_train_save_dir_overwrite(
     config = DummyTrainConfig(
         labels=labels_absolute_path,
         data_dir=TEST_VIDEOS_DIR,
-        model_name="dummy",
         checkpoint=dummy_checkpoint,
         save_dir=tmp_path / "my_model",
-        skip_load_validation=True,
         overwrite=True,
-        max_epochs=1,
-        batch_size=1,
-        auto_lr_find=False,
         num_workers=1,
     )
 

zamba/models/config.py

@@ -7,6 +7,7 @@
 import appdirs
 import ffmpeg
 from loguru import logger
+import numpy as np
 import pandas as pd
 from pydantic import BaseModel
 from pydantic import DirectoryPath, FilePath, validator, root_validator
@@ -495,19 +496,28 @@ def preprocess_labels(cls, values):
                     labels["site"], proportions=values["split_proportions"]
                 )
             else:
-                logger.info(
-                    "No 'site' column found so videos will be randomly allocated to splits."
-                )
                 # otherwise randomly allocate
-                random.seed(SPLIT_SEED)
-                labels["split"] = random.choices(
-                    list(values["split_proportions"].keys()),
-                    weights=list(values["split_proportions"].values()),
-                    k=len(labels),
+                logger.info(
+                    "No 'site' column found so videos will be randomly allocated using split proportions."
                 )
 
+                expected_labels = [k for k, v in values["split_proportions"].items() if v > 0]
+                labels["split"] = ""
+                seed = SPLIT_SEED
+
+                while len(np.setdiff1d(expected_labels, labels.split.unique())):
+
+                    random.seed(seed)
+                    labels["split"] = random.choices(
+                        list(values["split_proportions"].keys()),
+                        weights=list(values["split_proportions"].values()),
+                        k=len(labels),
+                    )
+
+                    seed += 1
+
                 logger.info(
-                    f"Writing out split information to {values['save_dir'] / 'splits.csv'}."
+                    f"Writing out split information to {values['save_dir'] / 'splits.csv'}. Used random seed {seed}."
                 )
 
                 # create the directory to save if we need to.

zamba/models/model_manager.py

@@ -245,8 +245,12 @@ def train_model(
     model_checkpoint = ModelCheckpoint(
         dirpath=logging_and_save_dir,
         filename=train_config.model_name,
-        monitor=train_config.early_stopping_config.monitor,
-        mode=train_config.early_stopping_config.mode,
+        monitor=train_config.early_stopping_config.monitor
+        if train_config.early_stopping_config is not None
+        else None,
+        mode=train_config.early_stopping_config.mode
+        if train_config.early_stopping_config is not None
+        else "min",
     )
 
     callbacks = [model_checkpoint]