From dd12098603ac415904b9a8d512889deb995a8391 Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Thu, 7 Jan 2021 20:54:44 +0800
Subject: Train different models in different conditions

---
 train.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

(limited to 'train.py')

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']
+    )
-- 
cgit v1.2.3