wehub-resource-sync
494 lines
16 KiB
Python
494 lines
16 KiB
Python
"""
|
|
Based on rllab's logger.
|
|
|
|
https://github.com/rll/rllab
|
|
"""
|
|
from enum import Enum
|
|
from contextlib import contextmanager
|
|
import numpy as np
|
|
import os
|
|
import os.path as osp
|
|
import sys
|
|
import datetime
|
|
import dateutil.tz
|
|
import csv
|
|
import json
|
|
import pickle
|
|
import errno
|
|
from collections import OrderedDict
|
|
from numbers import Number
|
|
import os
|
|
|
|
from tabulate import tabulate
|
|
import dateutil.tz
|
|
import os.path as osp
|
|
|
|
def dict_to_safe_json(d):
|
|
"""
|
|
Convert each value in the dictionary into a JSON'able primitive.
|
|
:param d:
|
|
:return:
|
|
"""
|
|
new_d = {}
|
|
for key, item in d.items():
|
|
if safe_json(item):
|
|
new_d[key] = item
|
|
else:
|
|
if isinstance(item, dict):
|
|
new_d[key] = dict_to_safe_json(item)
|
|
else:
|
|
new_d[key] = str(item)
|
|
return new_d
|
|
|
|
|
|
def safe_json(data):
|
|
if data is None:
|
|
return True
|
|
elif isinstance(data, (bool, int, float)):
|
|
return True
|
|
elif isinstance(data, (tuple, list)):
|
|
return all(safe_json(x) for x in data)
|
|
elif isinstance(data, dict):
|
|
return all(isinstance(k, str) and safe_json(v) for k, v in data.items())
|
|
return False
|
|
|
|
def create_exp_name(exp_prefix, exp_id=0, seed=0):
|
|
"""
|
|
Create a semi-unique experiment name that has a timestamp
|
|
:param exp_prefix:
|
|
:param exp_id:
|
|
:return:
|
|
"""
|
|
now = datetime.datetime.now(dateutil.tz.tzlocal())
|
|
timestamp = now.strftime('%Y_%m_%d_%H_%M_%S')
|
|
return "%s_%s_%04d--s-%d" % (exp_prefix, timestamp, exp_id, seed)
|
|
|
|
def create_log_dir(
|
|
exp_prefix,
|
|
exp_id=0,
|
|
seed=0,
|
|
base_log_dir=None,
|
|
include_exp_prefix_sub_dir=True,
|
|
):
|
|
"""
|
|
Creates and returns a unique log directory.
|
|
:param exp_prefix: All experiments with this prefix will have log
|
|
directories be under this directory.
|
|
:param exp_id: The number of the specific experiment run within this
|
|
experiment.
|
|
:param base_log_dir: The directory where all log should be saved.
|
|
:return:
|
|
"""
|
|
exp_name = create_exp_name(exp_prefix, exp_id=exp_id,
|
|
seed=seed)
|
|
if base_log_dir is None:
|
|
base_log_dir = './data'
|
|
if include_exp_prefix_sub_dir:
|
|
log_dir = osp.join(base_log_dir, exp_prefix.replace("_", "-"), exp_name)
|
|
else:
|
|
log_dir = osp.join(base_log_dir, exp_name)
|
|
if osp.exists(log_dir):
|
|
print("WARNING: Log directory already exists {}".format(log_dir), flush=True)
|
|
os.makedirs(log_dir, exist_ok=True)
|
|
return log_dir
|
|
|
|
|
|
def setup_logger(
|
|
exp_prefix="default",
|
|
variant=None,
|
|
text_log_file="debug.log",
|
|
variant_log_file="variant.json",
|
|
tabular_log_file="progress.csv",
|
|
snapshot_mode="last",
|
|
snapshot_gap=1,
|
|
log_tabular_only=False,
|
|
log_dir=None,
|
|
git_infos=None,
|
|
script_name=None,
|
|
**create_log_dir_kwargs
|
|
):
|
|
"""
|
|
Set up logger to have some reasonable default settings.
|
|
Will save log output to
|
|
based_log_dir/exp_prefix/exp_name.
|
|
exp_name will be auto-generated to be unique.
|
|
If log_dir is specified, then that directory is used as the output dir.
|
|
:param exp_prefix: The sub-directory for this specific experiment.
|
|
:param variant:
|
|
:param text_log_file:
|
|
:param variant_log_file:
|
|
:param tabular_log_file:
|
|
:param snapshot_mode:
|
|
:param log_tabular_only:
|
|
:param snapshot_gap:
|
|
:param log_dir:
|
|
:param git_infos:
|
|
:param script_name: If set, save the script name to this.
|
|
:return:
|
|
"""
|
|
first_time = log_dir is None
|
|
if first_time:
|
|
log_dir = create_log_dir(exp_prefix, **create_log_dir_kwargs)
|
|
|
|
if variant is not None:
|
|
logger.log("Variant:")
|
|
logger.log(json.dumps(dict_to_safe_json(variant), indent=2))
|
|
variant_log_path = osp.join(log_dir, variant_log_file)
|
|
logger.log_variant(variant_log_path, variant)
|
|
|
|
tabular_log_path = osp.join(log_dir, tabular_log_file)
|
|
text_log_path = osp.join(log_dir, text_log_file)
|
|
|
|
logger.add_text_output(text_log_path)
|
|
if first_time:
|
|
logger.add_tabular_output(tabular_log_path)
|
|
else:
|
|
logger._add_output(tabular_log_path, logger._tabular_outputs,
|
|
logger._tabular_fds, mode='a')
|
|
for tabular_fd in logger._tabular_fds:
|
|
logger._tabular_header_written.add(tabular_fd)
|
|
logger.set_snapshot_dir(log_dir)
|
|
logger.set_snapshot_mode(snapshot_mode)
|
|
logger.set_snapshot_gap(snapshot_gap)
|
|
logger.set_log_tabular_only(log_tabular_only)
|
|
exp_name = log_dir.split("/")[-1]
|
|
logger.push_prefix("[%s] " % exp_name)
|
|
|
|
if script_name is not None:
|
|
with open(osp.join(log_dir, "script_name.txt"), "w") as f:
|
|
f.write(script_name)
|
|
return log_dir
|
|
|
|
|
|
def create_stats_ordered_dict(
|
|
name,
|
|
data,
|
|
stat_prefix=None,
|
|
always_show_all_stats=True,
|
|
exclude_max_min=False,
|
|
):
|
|
if stat_prefix is not None:
|
|
name = "{}{}".format(stat_prefix, name)
|
|
if isinstance(data, Number):
|
|
return OrderedDict({name: data})
|
|
|
|
if len(data) == 0:
|
|
return OrderedDict()
|
|
|
|
if isinstance(data, tuple):
|
|
ordered_dict = OrderedDict()
|
|
for number, d in enumerate(data):
|
|
sub_dict = create_stats_ordered_dict(
|
|
"{0}_{1}".format(name, number),
|
|
d,
|
|
)
|
|
ordered_dict.update(sub_dict)
|
|
return ordered_dict
|
|
|
|
if isinstance(data, list):
|
|
try:
|
|
iter(data[0])
|
|
except TypeError:
|
|
pass
|
|
else:
|
|
data = np.concatenate(data)
|
|
|
|
if (isinstance(data, np.ndarray) and data.size == 1
|
|
and not always_show_all_stats):
|
|
return OrderedDict({name: float(data)})
|
|
|
|
stats = OrderedDict([
|
|
(name + ' Mean', np.mean(data)),
|
|
(name + ' Std', np.std(data)),
|
|
])
|
|
if not exclude_max_min:
|
|
stats[name + ' Max'] = np.max(data)
|
|
stats[name + ' Min'] = np.min(data)
|
|
return stats
|
|
|
|
|
|
class TerminalTablePrinter(object):
|
|
def __init__(self):
|
|
self.headers = None
|
|
self.tabulars = []
|
|
|
|
def print_tabular(self, new_tabular):
|
|
if self.headers is None:
|
|
self.headers = [x[0] for x in new_tabular]
|
|
else:
|
|
assert len(self.headers) == len(new_tabular)
|
|
self.tabulars.append([x[1] for x in new_tabular])
|
|
self.refresh()
|
|
|
|
def refresh(self):
|
|
import os
|
|
rows, columns = os.popen('stty size', 'r').read().split()
|
|
tabulars = self.tabulars[-(int(rows) - 3):]
|
|
sys.stdout.write("\x1b[2J\x1b[H")
|
|
sys.stdout.write(tabulate(tabulars, self.headers))
|
|
sys.stdout.write("\n")
|
|
|
|
|
|
class MyEncoder(json.JSONEncoder):
|
|
def default(self, o):
|
|
if isinstance(o, type):
|
|
return {'$class': o.__module__ + "." + o.__name__}
|
|
elif isinstance(o, Enum):
|
|
return {
|
|
'$enum': o.__module__ + "." + o.__class__.__name__ + '.' + o.name
|
|
}
|
|
elif callable(o):
|
|
return {
|
|
'$function': o.__module__ + "." + o.__name__
|
|
}
|
|
return json.JSONEncoder.default(self, o)
|
|
|
|
|
|
def mkdir_p(path):
|
|
try:
|
|
os.makedirs(path)
|
|
except OSError as exc: # Python >2.5
|
|
if exc.errno == errno.EEXIST and os.path.isdir(path):
|
|
pass
|
|
else:
|
|
raise
|
|
|
|
|
|
class Logger(object):
|
|
def __init__(self):
|
|
self._prefixes = []
|
|
self._prefix_str = ''
|
|
|
|
self._tabular_prefixes = []
|
|
self._tabular_prefix_str = ''
|
|
|
|
self._tabular = []
|
|
|
|
self._text_outputs = []
|
|
self._tabular_outputs = []
|
|
|
|
self._text_fds = {}
|
|
self._tabular_fds = {}
|
|
self._tabular_header_written = set()
|
|
|
|
self._snapshot_dir = None
|
|
self._snapshot_mode = 'all'
|
|
self._snapshot_gap = 1
|
|
|
|
self._log_tabular_only = False
|
|
self._header_printed = False
|
|
self.table_printer = TerminalTablePrinter()
|
|
|
|
def reset(self):
|
|
self.__init__()
|
|
|
|
def _add_output(self, file_name, arr, fds, mode='a'):
|
|
if file_name not in arr:
|
|
mkdir_p(os.path.dirname(file_name))
|
|
arr.append(file_name)
|
|
fds[file_name] = open(file_name, mode)
|
|
|
|
def _remove_output(self, file_name, arr, fds):
|
|
if file_name in arr:
|
|
fds[file_name].close()
|
|
del fds[file_name]
|
|
arr.remove(file_name)
|
|
|
|
def push_prefix(self, prefix):
|
|
self._prefixes.append(prefix)
|
|
self._prefix_str = ''.join(self._prefixes)
|
|
|
|
def add_text_output(self, file_name):
|
|
self._add_output(file_name, self._text_outputs, self._text_fds,
|
|
mode='a')
|
|
|
|
def remove_text_output(self, file_name):
|
|
self._remove_output(file_name, self._text_outputs, self._text_fds)
|
|
|
|
def add_tabular_output(self, file_name, relative_to_snapshot_dir=False):
|
|
if relative_to_snapshot_dir:
|
|
file_name = osp.join(self._snapshot_dir, file_name)
|
|
self._add_output(file_name, self._tabular_outputs, self._tabular_fds,
|
|
mode='w')
|
|
|
|
def remove_tabular_output(self, file_name, relative_to_snapshot_dir=False):
|
|
if relative_to_snapshot_dir:
|
|
file_name = osp.join(self._snapshot_dir, file_name)
|
|
if self._tabular_fds[file_name] in self._tabular_header_written:
|
|
self._tabular_header_written.remove(self._tabular_fds[file_name])
|
|
self._remove_output(file_name, self._tabular_outputs, self._tabular_fds)
|
|
|
|
def set_snapshot_dir(self, dir_name):
|
|
self._snapshot_dir = dir_name
|
|
|
|
def get_snapshot_dir(self, ):
|
|
return self._snapshot_dir
|
|
|
|
def get_snapshot_mode(self, ):
|
|
return self._snapshot_mode
|
|
|
|
def set_snapshot_mode(self, mode):
|
|
self._snapshot_mode = mode
|
|
|
|
def get_snapshot_gap(self, ):
|
|
return self._snapshot_gap
|
|
|
|
def set_snapshot_gap(self, gap):
|
|
self._snapshot_gap = gap
|
|
|
|
def set_log_tabular_only(self, log_tabular_only):
|
|
self._log_tabular_only = log_tabular_only
|
|
|
|
def get_log_tabular_only(self, ):
|
|
return self._log_tabular_only
|
|
|
|
def log(self, s, with_prefix=True, with_timestamp=True):
|
|
out = s
|
|
if with_prefix:
|
|
out = self._prefix_str + out
|
|
if with_timestamp:
|
|
now = datetime.datetime.now(dateutil.tz.tzlocal())
|
|
timestamp = now.strftime('%y-%m-%d.%H:%M') # :%S
|
|
out = "%s|%s" % (timestamp, out)
|
|
if not self._log_tabular_only:
|
|
# Also log to stdout
|
|
print(out, flush=True)
|
|
for fd in list(self._text_fds.values()):
|
|
fd.write(out + '\n')
|
|
fd.flush()
|
|
sys.stdout.flush()
|
|
|
|
def record_tabular(self, key, val):
|
|
self._tabular.append((self._tabular_prefix_str + str(key), str(val)))
|
|
|
|
def record_dict(self, d, prefix=None):
|
|
if prefix is not None:
|
|
self.push_tabular_prefix(prefix)
|
|
for k, v in d.items():
|
|
self.record_tabular(k, v)
|
|
if prefix is not None:
|
|
self.pop_tabular_prefix()
|
|
|
|
def push_tabular_prefix(self, key):
|
|
self._tabular_prefixes.append(key)
|
|
self._tabular_prefix_str = ''.join(self._tabular_prefixes)
|
|
|
|
def pop_tabular_prefix(self, ):
|
|
del self._tabular_prefixes[-1]
|
|
self._tabular_prefix_str = ''.join(self._tabular_prefixes)
|
|
|
|
def save_extra_data(self, data, file_name='extra_data.pkl', mode='joblib'):
|
|
"""
|
|
Data saved here will always override the last entry
|
|
|
|
:param data: Something pickle'able.
|
|
"""
|
|
file_name = osp.join(self._snapshot_dir, file_name)
|
|
if mode == 'joblib':
|
|
import joblib
|
|
joblib.dump(data, file_name, compress=3)
|
|
elif mode == 'pickle':
|
|
pickle.dump(data, open(file_name, "wb"))
|
|
else:
|
|
raise ValueError("Invalid mode: {}".format(mode))
|
|
return file_name
|
|
|
|
def get_table_dict(self, ):
|
|
return dict(self._tabular)
|
|
|
|
def get_table_key_set(self, ):
|
|
return set(key for key, value in self._tabular)
|
|
|
|
@contextmanager
|
|
def prefix(self, key):
|
|
self.push_prefix(key)
|
|
try:
|
|
yield
|
|
finally:
|
|
self.pop_prefix()
|
|
|
|
@contextmanager
|
|
def tabular_prefix(self, key):
|
|
self.push_tabular_prefix(key)
|
|
yield
|
|
self.pop_tabular_prefix()
|
|
|
|
def log_variant(self, log_file, variant_data):
|
|
mkdir_p(os.path.dirname(log_file))
|
|
with open(log_file, "w") as f:
|
|
json.dump(variant_data, f, indent=2, sort_keys=True, cls=MyEncoder)
|
|
|
|
def record_tabular_misc_stat(self, key, values, placement='back'):
|
|
if placement == 'front':
|
|
prefix = ""
|
|
suffix = key
|
|
else:
|
|
prefix = key
|
|
suffix = ""
|
|
if len(values) > 0:
|
|
self.record_tabular(prefix + "Average" + suffix, np.average(values))
|
|
self.record_tabular(prefix + "Std" + suffix, np.std(values))
|
|
self.record_tabular(prefix + "Median" + suffix, np.median(values))
|
|
self.record_tabular(prefix + "Min" + suffix, np.min(values))
|
|
self.record_tabular(prefix + "Max" + suffix, np.max(values))
|
|
else:
|
|
self.record_tabular(prefix + "Average" + suffix, np.nan)
|
|
self.record_tabular(prefix + "Std" + suffix, np.nan)
|
|
self.record_tabular(prefix + "Median" + suffix, np.nan)
|
|
self.record_tabular(prefix + "Min" + suffix, np.nan)
|
|
self.record_tabular(prefix + "Max" + suffix, np.nan)
|
|
|
|
def dump_tabular(self, *args, **kwargs):
|
|
wh = kwargs.pop("write_header", None)
|
|
if len(self._tabular) > 0:
|
|
if self._log_tabular_only:
|
|
self.table_printer.print_tabular(self._tabular)
|
|
else:
|
|
for line in tabulate(self._tabular).split('\n'):
|
|
self.log(line, *args, **kwargs)
|
|
tabular_dict = dict(self._tabular)
|
|
# Also write to the csv files
|
|
# This assumes that the keys in each iteration won't change!
|
|
for tabular_fd in list(self._tabular_fds.values()):
|
|
writer = csv.DictWriter(tabular_fd,
|
|
fieldnames=list(tabular_dict.keys()))
|
|
if wh or (
|
|
wh is None and tabular_fd not in self._tabular_header_written):
|
|
writer.writeheader()
|
|
self._tabular_header_written.add(tabular_fd)
|
|
writer.writerow(tabular_dict)
|
|
tabular_fd.flush()
|
|
del self._tabular[:]
|
|
|
|
def pop_prefix(self, ):
|
|
del self._prefixes[-1]
|
|
self._prefix_str = ''.join(self._prefixes)
|
|
|
|
def save_itr_params(self, itr, params):
|
|
if self._snapshot_dir:
|
|
if self._snapshot_mode == 'all':
|
|
file_name = osp.join(self._snapshot_dir, 'itr_%d.pkl' % itr)
|
|
pickle.dump(params, open(file_name, "wb"))
|
|
elif self._snapshot_mode == 'last':
|
|
# override previous params
|
|
file_name = osp.join(self._snapshot_dir, 'params.pkl')
|
|
pickle.dump(params, open(file_name, "wb"))
|
|
elif self._snapshot_mode == "gap":
|
|
if itr % self._snapshot_gap == 0:
|
|
file_name = osp.join(self._snapshot_dir, 'itr_%d.pkl' % itr)
|
|
pickle.dump(params, open(file_name, "wb"))
|
|
elif self._snapshot_mode == "gap_and_last":
|
|
if itr % self._snapshot_gap == 0:
|
|
file_name = osp.join(self._snapshot_dir, 'itr_%d.pkl' % itr)
|
|
pickle.dump(params, open(file_name, "wb"))
|
|
file_name = osp.join(self._snapshot_dir, 'params.pkl')
|
|
pickle.dump(params, open(file_name, "wb"))
|
|
elif self._snapshot_mode == 'none':
|
|
pass
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
|
|
logger = Logger()
|
|
|