diff options
| -rw-r--r-- | models/model.py | 6 | ||||
| -rw-r--r-- | train.py | 14 |
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) @@ -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'] + ) |