summaryrefslogtreecommitdiff
path: root/train.py
blob: e845d6f54af5a859e56ee24dd4e13e9e1fbbe2f2 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
from config import config
from models import Model
from utils.dataset import ClipConditions

model = Model(config['system'], config['model'], config['hyperparameter'])

# 3 models for different conditions
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'})},
}

model.fit_all(config['dataset'], dataset_selectors, config['dataloader'])