chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,55 @@
|
||||
# Copyright 2022 Twitter, Inc and Zhendong Wang.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
class Data_Sampler(object):
|
||||
def __init__(self, data, device, reward_tune='no'):
|
||||
self.state = torch.from_numpy(data['observations']).float()
|
||||
self.action = torch.from_numpy(data['actions']).float()
|
||||
self.next_state = torch.from_numpy(data['next_observations']).float()
|
||||
reward = torch.from_numpy(data['rewards']).reshape(-1, 1).float()
|
||||
self.not_done = 1. - torch.from_numpy(data['dones']).reshape(-1, 1).float()
|
||||
|
||||
self.size = self.state.shape[0]
|
||||
self.state_dim = self.state.shape[1]
|
||||
self.action_dim = self.action.shape[1]
|
||||
self.device = device
|
||||
|
||||
# 保留原 reward 调整逻辑
|
||||
if reward_tune == 'normalize':
|
||||
reward = (reward - reward.mean()) / reward.std()
|
||||
elif reward_tune == 'iql_antmaze':
|
||||
reward = reward - 1.0
|
||||
elif reward_tune == 'iql_locomotion':
|
||||
reward = iql_normalize(reward, self.not_done)
|
||||
elif reward_tune == 'cql_antmaze':
|
||||
reward = (reward - 0.5) * 4.0
|
||||
elif reward_tune == 'antmaze':
|
||||
reward = (reward - 0.25) * 2.0
|
||||
|
||||
self.reward = reward
|
||||
|
||||
def sample(self, batch_size):
|
||||
ind = torch.randint(0, self.size, size=(batch_size,))
|
||||
return (
|
||||
self.state[ind].to(self.device),
|
||||
self.action[ind].to(self.device),
|
||||
self.next_state[ind].to(self.device),
|
||||
self.reward[ind].to(self.device),
|
||||
self.not_done[ind].to(self.device)
|
||||
)
|
||||
|
||||
def iql_normalize(reward, not_done):
|
||||
trajs_rt = []
|
||||
episode_return = 0.0
|
||||
for i in range(len(reward)):
|
||||
episode_return += reward[i]
|
||||
if not not_done[i]:
|
||||
trajs_rt.append(episode_return)
|
||||
episode_return = 0.0
|
||||
rt_max, rt_min = torch.max(torch.tensor(trajs_rt)), torch.min(torch.tensor(trajs_rt))
|
||||
reward /= (rt_max - rt_min)
|
||||
reward *= 1000.
|
||||
return reward
|
||||
@@ -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()
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
# Copyright 2022 Twitter, Inc and Zhendong Wang.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def soft_update_from_to(source, target, tau):
|
||||
for target_param, param in zip(target.parameters(), source.parameters()):
|
||||
target_param.data.copy_(
|
||||
target_param.data * (1.0 - tau) + param.data * tau
|
||||
)
|
||||
|
||||
|
||||
def copy_model_params_from_to(source, target):
|
||||
for target_param, param in zip(target.parameters(), source.parameters()):
|
||||
target_param.data.copy_(param.data)
|
||||
|
||||
|
||||
def fanin_init(tensor, scale=1):
|
||||
size = tensor.size()
|
||||
if len(size) == 2:
|
||||
fan_in = size[0]
|
||||
elif len(size) > 2:
|
||||
fan_in = np.prod(size[1:])
|
||||
else:
|
||||
raise Exception("Shape must be have dimension at least 2.")
|
||||
bound = scale / np.sqrt(fan_in)
|
||||
return tensor.data.uniform_(-bound, bound)
|
||||
|
||||
|
||||
def orthogonal_init(tensor, gain=0.01):
|
||||
torch.nn.init.orthogonal_(tensor, gain=gain)
|
||||
|
||||
|
||||
def fanin_init_weights_like(tensor):
|
||||
size = tensor.size()
|
||||
if len(size) == 2:
|
||||
fan_in = size[0]
|
||||
elif len(size) > 2:
|
||||
fan_in = np.prod(size[1:])
|
||||
else:
|
||||
raise Exception("Shape must be have dimension at least 2.")
|
||||
bound = 1. / np.sqrt(fan_in)
|
||||
new_tensor = torch.FloatTensor(tensor.size())
|
||||
new_tensor.uniform_(-bound, bound)
|
||||
return new_tensor
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
# Copyright 2022 Twitter, Inc and Zhendong Wang.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import time
|
||||
import math
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def print_banner(s, separator="-", num_star=60):
|
||||
print(separator * num_star, flush=True)
|
||||
print(s, flush=True)
|
||||
print(separator * num_star, flush=True)
|
||||
|
||||
|
||||
class Progress:
|
||||
|
||||
def __init__(self, total, name='Progress', ncol=3, max_length=20, indent=0, line_width=100, speed_update_freq=100):
|
||||
self.total = total
|
||||
self.name = name
|
||||
self.ncol = ncol
|
||||
self.max_length = max_length
|
||||
self.indent = indent
|
||||
self.line_width = line_width
|
||||
self._speed_update_freq = speed_update_freq
|
||||
|
||||
self._step = 0
|
||||
self._prev_line = '\033[F'
|
||||
self._clear_line = ' ' * self.line_width
|
||||
|
||||
self._pbar_size = self.ncol * self.max_length
|
||||
self._complete_pbar = '#' * self._pbar_size
|
||||
self._incomplete_pbar = ' ' * self._pbar_size
|
||||
|
||||
self.lines = ['']
|
||||
self.fraction = '{} / {}'.format(0, self.total)
|
||||
|
||||
self.resume()
|
||||
|
||||
def update(self, description, n=1):
|
||||
self._step += n
|
||||
if self._step % self._speed_update_freq == 0:
|
||||
self._time0 = time.time()
|
||||
self._step0 = self._step
|
||||
self.set_description(description)
|
||||
|
||||
def resume(self):
|
||||
self._skip_lines = 1
|
||||
print('\n', end='')
|
||||
self._time0 = time.time()
|
||||
self._step0 = self._step
|
||||
|
||||
def pause(self):
|
||||
self._clear()
|
||||
self._skip_lines = 1
|
||||
|
||||
def set_description(self, params=[]):
|
||||
|
||||
if type(params) == dict:
|
||||
params = sorted([
|
||||
(key, val)
|
||||
for key, val in params.items()
|
||||
])
|
||||
|
||||
############
|
||||
# Position #
|
||||
############
|
||||
self._clear()
|
||||
|
||||
###########
|
||||
# Percent #
|
||||
###########
|
||||
percent, fraction = self._format_percent(self._step, self.total)
|
||||
self.fraction = fraction
|
||||
|
||||
#########
|
||||
# Speed #
|
||||
#########
|
||||
speed = self._format_speed(self._step)
|
||||
|
||||
##########
|
||||
# Params #
|
||||
##########
|
||||
num_params = len(params)
|
||||
nrow = math.ceil(num_params / self.ncol)
|
||||
params_split = self._chunk(params, self.ncol)
|
||||
params_string, lines = self._format(params_split)
|
||||
self.lines = lines
|
||||
|
||||
description = '{} | {}{}'.format(percent, speed, params_string)
|
||||
print(description)
|
||||
self._skip_lines = nrow + 1
|
||||
|
||||
def append_description(self, descr):
|
||||
self.lines.append(descr)
|
||||
|
||||
def _clear(self):
|
||||
position = self._prev_line * self._skip_lines
|
||||
empty = '\n'.join([self._clear_line for _ in range(self._skip_lines)])
|
||||
print(position, end='')
|
||||
print(empty)
|
||||
print(position, end='')
|
||||
|
||||
def _format_percent(self, n, total):
|
||||
if total:
|
||||
percent = n / float(total)
|
||||
|
||||
complete_entries = int(percent * self._pbar_size)
|
||||
incomplete_entries = self._pbar_size - complete_entries
|
||||
|
||||
pbar = self._complete_pbar[:complete_entries] + self._incomplete_pbar[:incomplete_entries]
|
||||
fraction = '{} / {}'.format(n, total)
|
||||
string = '{} [{}] {:3d}%'.format(fraction, pbar, int(percent * 100))
|
||||
else:
|
||||
fraction = '{}'.format(n)
|
||||
string = '{} iterations'.format(n)
|
||||
return string, fraction
|
||||
|
||||
def _format_speed(self, n):
|
||||
num_steps = n - self._step0
|
||||
t = time.time() - self._time0
|
||||
speed = num_steps / t
|
||||
string = '{:.1f} Hz'.format(speed)
|
||||
if num_steps > 0:
|
||||
self._speed = string
|
||||
return string
|
||||
|
||||
def _chunk(self, l, n):
|
||||
return [l[i:i + n] for i in range(0, len(l), n)]
|
||||
|
||||
def _format(self, chunks):
|
||||
lines = [self._format_chunk(chunk) for chunk in chunks]
|
||||
lines.insert(0, '')
|
||||
padding = '\n' + ' ' * self.indent
|
||||
string = padding.join(lines)
|
||||
return string, lines
|
||||
|
||||
def _format_chunk(self, chunk):
|
||||
line = ' | '.join([self._format_param(param) for param in chunk])
|
||||
return line
|
||||
|
||||
def _format_param(self, param):
|
||||
k, v = param
|
||||
return '{} : {}'.format(k, v)[:self.max_length]
|
||||
|
||||
def stamp(self):
|
||||
if self.lines != ['']:
|
||||
params = ' | '.join(self.lines)
|
||||
string = '[ {} ] {}{} | {}'.format(self.name, self.fraction, params, self._speed)
|
||||
self._clear()
|
||||
print(string, end='\n')
|
||||
self._skip_lines = 1
|
||||
else:
|
||||
self._clear()
|
||||
self._skip_lines = 0
|
||||
|
||||
def close(self):
|
||||
self.pause()
|
||||
|
||||
|
||||
class Silent:
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def __getattr__(self, attr):
|
||||
return lambda *args: None
|
||||
|
||||
|
||||
class EarlyStopping(object):
|
||||
def __init__(self, tolerance=5, min_delta=0):
|
||||
self.tolerance = tolerance
|
||||
self.min_delta = min_delta
|
||||
self.counter = 0
|
||||
self.early_stop = False
|
||||
|
||||
def __call__(self, train_loss, validation_loss):
|
||||
if (validation_loss - train_loss) > self.min_delta:
|
||||
self.counter += 1
|
||||
if self.counter >= self.tolerance:
|
||||
return True
|
||||
else:
|
||||
self.counter = 0
|
||||
return False
|
||||
Reference in New Issue
Block a user