blob: cdb2fb04d3fd04fa92f8dfd65023aa125d8d5418 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
|
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)
if CUDA_VISIBLE_DEVICES:
os.environ['CUDA_VISIBLE_DEVICES'] = CUDA_VISIBLE_DEVICES
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'})},
}
for selector in dataset_selectors.values():
model.fit(
dict(**config['dataset'], **{'selector': selector}),
config['dataloader']
)
|
