def dump_tabular(self):
"""
Write all of the diagnostics from the current iteration.
Writes both to stdout, and to the output file.
"""
if proc_id() == 0:
vals = []
key_lens = [len(key) for key in self.log_headers]
max_key_len = max(15, max(key_lens))
keystr = '%' + '%d' % max_key_len
fmt = "| " + keystr + "s | %15s |"
n_slashes = 22 + max_key_len
print("-" * n_slashes)
for key in self.log_headers:
val = self.log_current_row.get(key, "")
valstr = "%8.3g" % val if hasattr(val, "__float__") else val
print(fmt % (key, valstr))
vals.append(val)
print("-" * n_slashes, flush=True)
if self.output_file is not None:
if self.first_row:
self.output_file.write("\t".join(self.log_headers) + "\n")
self.output_file.write("\t".join(map(str, vals)) + "\n")
self.output_file.flush()
self.log_current_row.clear()
self.first_row = False
def save_state(self, state_dict, itr=None):
"""
Saves the state of an experiment.
To be clear: this is about saving *state*, not logging diagnostics.
All diagnostic logging is separate from this function. This function
will save whatever is in ``state_dict``---usually just a copy of the
environment---and the most recent parameters for the model you
previously set up saving for with ``setup_tf_saver``.
Call with any frequency you prefer. If you only want to maintain a
single state and overwrite it at each call with the most recent
version, leave ``itr=None``. If you want to keep all of the states you
save, provide unique (increasing) values for 'itr'.
Args:
state_dict (dict): Dictionary containing essential elements to
describe the current state of training.
itr: An int, or None. Current iteration of training.
"""
if proc_id() == 0:
fname = 'vars.pkl' if itr is None else 'vars%d.pkl' % itr
try:
joblib.dump(state_dict, osp.join(self.output_dir, fname))
except:
self.log('Warning: could not pickle state_dict.', color='red')
if hasattr(self, 'tf_saver_elements'):
self._tf_simple_save(itr)
if hasattr(self, 'pytorch_saver_elements'):
self._pytorch_simple_save(itr)
def save_config(self, config):
"""
Log an experiment configuration.
Call this once at the top of your experiment, passing in all important
config vars as a dict. This will serialize the config to JSON, while
handling anything which can't be serialized in a graceful way (writing
as informative a string as possible).
Example use:
.. code-block:: python
logger = EpochLogger(**logger_kwargs)
logger.save_config(locals())
"""
config_json = convert_json(config)
if self.exp_name is not None:
config_json['exp_name'] = self.exp_name
if proc_id() == 0:
output = json.dumps(config_json,
separators=(',', ':\t'),
indent=4,
sort_keys=True)
print(colorize('Saving config:\n', color='cyan', bold=True))
print(output)
with open(osp.join(self.output_dir, "config.json"), 'w') as out:
out.write(output)
def _tf_simple_save(self, itr=None):
"""
Uses simple_save to save a trained model, plus info to make it easy
to associated tensors to variables after restore.
"""
if proc_id() == 0:
assert hasattr(self, 'tf_saver_elements'), \
"First have to setup saving with self.setup_tf_saver"
fpath = 'tf1_save' + ('%d' % itr if itr is not None else '')
fpath = osp.join(self.output_dir, fpath)
if osp.exists(fpath):
# simple_save refuses to be useful if fpath already exists,
# so just delete fpath if it's there.
shutil.rmtree(fpath)
tf.saved_model.simple_save(export_dir=fpath,
**self.tf_saver_elements)
joblib.dump(self.tf_saver_info, osp.join(fpath, 'model_info.pkl'))
def __init__(self,
output_dir=None,
output_fname='progress.txt',
exp_name=None):
"""
Initialize a Logger.
Args:
output_dir (string): A directory for saving results to. If
``None``, defaults to a temp directory of the form
``/tmp/experiments/somerandomnumber``.
output_fname (string): Name for the tab-separated-value file
containing metrics logged throughout a training run.
Defaults to ``progress.txt``.
exp_name (string): Experiment name. If you run multiple training
runs and give them all the same ``exp_name``, the plotter
will know to group them. (Use case: if you run the same
hyperparameter configuration with multiple random seeds, you
should give them all the same ``exp_name``.)
"""
if proc_id() == 0:
self.output_dir = output_dir or "/tmp/experiments/%i" % int(
time.time())
if osp.exists(self.output_dir):
print(
"Warning: Log dir %s already exists! Storing info there anyway."
% self.output_dir)
else:
os.makedirs(self.output_dir)
self.output_file = open(osp.join(self.output_dir, output_fname),
'w')
atexit.register(self.output_file.close)
print(
colorize("Logging data to %s" % self.output_file.name,
'green',
bold=True))
else:
self.output_dir = None
self.output_file = None
self.first_row = True
self.log_headers = []
self.log_current_row = {}
self.exp_name = exp_name
def _pytorch_simple_save(self, itr=None):
"""
Saves the PyTorch model (or models).
"""
if proc_id() == 0:
assert hasattr(self, 'pytorch_saver_elements'), \
"First have to setup saving with self.setup_pytorch_saver"
fpath = 'pyt_save'
fpath = osp.join(self.output_dir, fpath)
fname = 'model' + ('%d' % itr if itr is not None else '') + '.pt'
fname = osp.join(fpath, fname)
os.makedirs(fpath, exist_ok=True)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
# We are using a non-recommended way of saving PyTorch models,
# by pickling whole objects (which are dependent on the exact
# directory structure at the time of saving) as opposed to
# just saving network weights. This works sufficiently well
# for the purposes of Spinning Up, but you may want to do
# something different for your personal PyTorch project.
# We use a catch_warnings() context to avoid the warnings about
# not being able to save the source code.
torch.save(self.pytorch_saver_elements, fname)
def log(self, msg, color='green'):
"""Print a colorized message to stdout."""
if proc_id() == 0:
print(colorize(msg, color, bold=True))