From d30cf2cb280e83e4a4abe1e9c2abdbba17d903a3 Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Sun, 10 Jan 2021 19:54:42 +0800 Subject: Make predict function transform samples different conditions in a single shot --- train.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) (limited to 'train.py') diff --git a/train.py b/train.py index d921839..cdb2fb0 100644 --- a/train.py +++ b/train.py @@ -12,12 +12,12 @@ if CUDA_VISIBLE_DEVICES: model = Model(config['system'], config['model'], config['hyperparameter']) # 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: +dataset_selectors = { + 'nm': {'conditions': ClipConditions({r'nm-0\d'})}, + 'bg': {'conditions': ClipConditions({r'nm-0\d', r'bg-0\d'})}, + 'cl': {'conditions': ClipConditions({r'nm-0\d', r'cl-0\d'})}, +} +for selector in dataset_selectors.values(): model.fit( dict(**config['dataset'], **{'selector': selector}), config['dataloader'] -- cgit v1.2.3