Train different models in different conditions

2 files changed, 17 insertions, 3 deletions

diff --git a/models/model.py b/models/model.py
index 1dc0f23..54f3441 100644
--- a/models/model.py
+++ b/models/model.py
@@ -220,13 +220,15 @@ class Model:
 
         return self._gen_sig(list(_config.values()))
 
-    def _gen_sig(self, values: Union[tuple, list, str, int, float]) -> str:
+    def _gen_sig(self, values: Union[tuple, list, set, str, int, float]) -> str:
         strings = []
         for v in values:
             if isinstance(v, str):
                 strings.append(v)
-            elif isinstance(v, (tuple, list)):
+            elif isinstance(v, (tuple, list, set)):
                 strings.append(self._gen_sig(v))
+            elif isinstance(v, dict):
+                strings.append(self._gen_sig(list(v.values())))
             else:
                 strings.append(str(v))
         return '_'.join(strings)
diff --git a/train.py b/train.py
index 17cd0f6..d921839 100644
--- a/train.py
+++ b/train.py
@@ -2,6 +2,7 @@ import os
 
 from config import config
 from models import Model
+from utils.dataset import ClipConditions
 
 # Set environment variable CUDA device(s)
 CUDA_VISIBLE_DEVICES = config['system'].get('CUDA_VISIBLE_DEVICES', None)
@@ -9,4 +10,15 @@ if CUDA_VISIBLE_DEVICES:
     os.environ['CUDA_VISIBLE_DEVICES'] = CUDA_VISIBLE_DEVICES
 
 model = Model(config['system'], config['model'], config['hyperparameter'])
-model.fit(config['dataset'], config['dataloader'])
+
+# 3 models for different conditions
+dataset_selectors = [
+    {'conditions': ClipConditions({r'nm-0\d'})},
+    {'conditions': ClipConditions({r'nm-0\d', r'bg-0\d'})},
+    {'conditions': ClipConditions({r'nm-0\d', r'cl-0\d'})},
+]
+for selector in dataset_selectors:
+    model.fit(
+        dict(**config['dataset'], **{'selector': selector}),
+        config['dataloader']
+    )