-
Notifications
You must be signed in to change notification settings - Fork 26
/
utils.py
44 lines (34 loc) · 1.21 KB
/
utils.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
import os, sys
import yaml
def get_result_dir(args):
'''
Get directory of result location (result/problem/subfolder/algo/seed/)
'''
top_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'result')
exp_name = '' if args.exp_name is None else '-' + args.exp_name
algo_name = args.algo + exp_name
result_dir = os.path.join(top_dir, args.problem, args.subfolder, algo_name, str(args.seed))
os.makedirs(result_dir, exist_ok=True)
return result_dir
def save_args(general_args, framework_args):
'''
Save arguments to yaml file
'''
all_args = {'general': vars(general_args)}
all_args.update(framework_args)
result_dir = get_result_dir(general_args)
args_path = os.path.join(result_dir, 'args.yml')
os.makedirs(os.path.dirname(args_path), exist_ok=True)
with open(args_path, 'w') as f:
yaml.dump(all_args, f, default_flow_style=False, sort_keys=False)
def setup_logger(args):
'''
Log to file if needed
'''
logger = None
if args.log_to_file:
result_dir = get_result_dir(args)
log_path = os.path.join(result_dir, 'log.txt')
logger = open(log_path, 'w')
sys.stdout = logger
return logger