/
run_lab.py
100 lines (86 loc) · 3.55 KB
/
run_lab.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
'''
The entry point of SLM Lab
Specify what to run in `config/experiments.json`
Then run `yarn start` or `python run_lab.py`
'''
import os
# NOTE increase if needed. Pytorch thread overusage https://github.com/pytorch/pytorch/issues/975
os.environ['OMP_NUM_THREADS'] = '1'
from slm_lab.experiment import analysis
from slm_lab.experiment.control import Session, Trial, Experiment
from slm_lab.experiment.monitor import InfoSpace
from slm_lab.lib import logger, util
from slm_lab.spec import spec_util
from xvfbwrapper import Xvfb
import sys
import torch.multiprocessing as mp
debug_modules = [
# 'algorithm',
]
debug_level = 'DEBUG'
logger.toggle_debug(debug_modules, debug_level)
def run_new_mode(spec_file, spec_name, lab_mode):
'''Run to generate new data with `search, train, dev`'''
spec = spec_util.get(spec_file, spec_name)
info_space = InfoSpace()
analysis.save_spec(spec, info_space, unit='experiment') # first save the new spec
if lab_mode == 'search':
info_space.tick('experiment')
Experiment(spec, info_space).run()
elif lab_mode.startswith('train'):
info_space.tick('trial')
Trial(spec, info_space).run()
elif lab_mode == 'dev':
spec = spec_util.override_dev_spec(spec)
info_space.tick('trial')
Trial(spec, info_space).run()
else:
raise ValueError('Unrecognizable lab_mode not of `search, train, dev`')
def run_old_mode(spec_file, spec_name, lab_mode):
'''Run using existing data with `enjoy, eval`. The eval mode is also what train mode's online eval runs in a subprocess via bash command'''
# reconstruct spec and info_space from existing data
lab_mode, prename = lab_mode.split('@')
predir, _, _, _, _, _ = util.prepath_split(spec_file)
prepath = f'{predir}/{prename}'
spec, info_space = util.prepath_to_spec_info_space(prepath)
# see InfoSpace def for more on these
info_space.ckpt = 'eval'
info_space.eval_model_prepath = prepath
# no info_space.tick() as they are reconstructed
if lab_mode == 'enjoy':
spec = spec_util.override_enjoy_spec(spec)
Session(spec, info_space).run()
elif lab_mode == 'eval':
spec = spec_util.override_eval_spec(spec)
Session(spec, info_space).run()
util.clear_periodic_ckpt(prepath) # cleanup after itself
analysis.analyze_eval_trial(spec, info_space, predir)
else:
raise ValueError('Unrecognizable lab_mode not of `enjoy, eval`')
def run_by_mode(spec_file, spec_name, lab_mode):
'''The main run lab function for all lab_modes'''
logger.info(f'Running lab in mode: {lab_mode}')
# '@' is reserved for 'enjoy@{prename}'
os.environ['lab_mode'] = lab_mode.split('@')[0]
if lab_mode in ('search', 'train', 'dev'):
run_new_mode(spec_file, spec_name, lab_mode)
else:
run_old_mode(spec_file, spec_name, lab_mode)
def main():
if len(sys.argv) > 1:
args = sys.argv[1:]
assert len(args) == 3, f'To use sys args, specify spec_file, spec_name, lab_mode'
run_by_mode(*args)
return
experiments = util.read('config/experiments.json')
for spec_file in experiments:
for spec_name, lab_mode in experiments[spec_file].items():
run_by_mode(spec_file, spec_name, lab_mode)
if __name__ == '__main__':
mp.set_start_method('spawn') # for distributed pytorch to work
if sys.platform == 'darwin':
# avoid xvfb for MacOS: https://github.com/nipy/nipype/issues/1400
main()
else:
with Xvfb() as xvfb: # safety context for headless machines
main()