memef4rmer's picture
Duplicate from jyseo/3DFuse
efe5745
# design inspiration from detectron2
from pathlib import Path
import json
import os
from contextlib import contextmanager
from .ticker import IntervalTicker
_CURRENT_STORAGE_STACK = []
def get_event_storage():
"""
Returns:
The :class:`EventStorage` object that's currently being used.
Throws an error if no :class:`EventStorage` is currently enabled.
"""
assert len(
_CURRENT_STORAGE_STACK
), "get_event_storage() has to be called inside a 'with EventStorage(...)' context!"
return _CURRENT_STORAGE_STACK[-1]
def read_lined_json(fname):
with Path(fname).open('r') as f:
for line in f:
item = json.loads(line)
yield item
def read_stats(dirname, key):
if dirname is None or not (fname := Path(dirname) / "history.json").is_file():
return [], []
stats = read_lined_json(fname)
stats = list(filter(lambda x: key in x, stats))
xs = [e['iter'] for e in stats]
ys = [e[key] for e in stats]
return xs, ys
class EventStorage():
def __init__(self, output_dir="./hotdog", start_iter=0, flush_period=60):
self.iter = start_iter
self.ticker = IntervalTicker(flush_period)
self.history = []
self._current_prefix = ""
self._init_curr_buffer_()
self.output_dir = output_dir
self.writable = False
def _open(self):
if self.writable:
output_dir = Path(self.output_dir)
if not output_dir.is_dir():
output_dir.mkdir(parents=True, exist_ok=True)
json_fname = output_dir / 'history.json'
self._file_handle = json_fname.open('a', encoding='utf8')
self.output_dir = output_dir # make sure it's a path object
def _init_curr_buffer_(self):
self.curr_buffer = {'iter': self.iter}
def step(self, flush=False):
self.history.append(self.curr_buffer)
on_flush_period = self.ticker.tick()
if flush or on_flush_period:
self.flush_history()
self.iter += 1
self._init_curr_buffer_()
def flush_history(self):
if self.writable:
for item in self.history:
line = json.dumps(item, sort_keys=True, ensure_ascii=False) + "\n"
self._file_handle.write(line)
self._file_handle.flush()
self.history = []
def full_key(self, key):
assert isinstance(key, str)
name = self._current_prefix + key
return name
def put(self, key, val):
key = self.full_key(key)
assert isinstance(val, (int, float, str))
if isinstance(val, float):
val = round(val, 3)
self.curr_buffer[key] = val
def put_scalars(self, **kwargs):
for k, v in kwargs.items():
self.put(k, v)
def put_artifact(self, key, ext,p, save_func):
if not self.writable:
return
p=p.replace(" ","_")
os.makedirs(self.output_dir / key, exist_ok=True)
fname = (self.output_dir / key / f"step_{self.iter}_{p}").with_suffix(ext)
fname = str(fname)
# must be called inside so that
# 1. the func is not executed if the metric is not writable
# 2. the key is only inserted if the func succeeds
save_func(fname)
self.put(key, fname)
return fname
def close(self):
self.flush_history()
if self.writable:
self._file_handle.close()
def get_last(self):
if len(self.history) > 0:
last = self.history[-1]
return last
def __enter__(self):
if len(_CURRENT_STORAGE_STACK) > 0:
parent = _CURRENT_STORAGE_STACK[-1]
root, dirname = parent.output_dir, self.output_dir
if root is not None and dirname is not None:
child_dir = parent.output_dir / f"{self.output_dir}_{parent.iter}"
self.output_dir = child_dir
parent.put(str(dirname), str(child_dir))
if self.output_dir is not None:
self.writable = True
self._open()
_CURRENT_STORAGE_STACK.append(self)
return self
def __exit__(self, exc_type, exc_val, exc_tb):
assert _CURRENT_STORAGE_STACK[-1] == self
_CURRENT_STORAGE_STACK.pop()
self.close()