chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,493 @@
|
||||
"""
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user