class TensorBoardOutputFormat(KVWriter):
"""
Dumps key/value pairs into TensorBoard's numeric format.
"""
def __init__(self, dir):
os.makedirs(dir, exist_ok=True)
self.dir = dir
self.step = 1
from torch.utils.tensorboard import SummaryWriter
self.writer = SummaryWriter(dir)
def writekvs(self, kvs):
def summary_val(k, v):
return k, float(v)
step = kvs.get('timesteps', self.step)
for k, v in kvs.items():
self.writer.add_scalar(*summary_val(k, v), step)
self.writer.flush()
self.step += 1
def close(self):
if self.writer:
self.writer.Close()
self.writer = None
class TensorBoardOutputFormat(KVWriter):
"""
Dumps key/value pairs into TensorBoard's numeric format.
"""
def __init__(self, dir):
# os.makedirs(dir, exist_ok=True)
# self.dir = dir
# self.step = 1
# prefix = 'events'
# path = osp.join(osp.abspath(dir), prefix)
# import tensorflow as tf
# from tensorflow.python import pywrap_tensorflow
# from tensorflow.core.util import event_pb2
# from tensorflow import compat
# self.tf = tf
# self.event_pb2 = event_pb2
# self.pywrap_tensorflow = pywrap_tensorflow
# self.writer = pywrap_tensorflow.EventsWriter(compat.as_bytes(path))
from torch.utils.tensorboard import SummaryWriter
self.writer = SummaryWriter(log_dir=dir)
def get_writer(self):
return self.writer
def writekvs(self, kvs):
step = int(kvs["steps"])
for k, v in kvs.items():
self.writer.add_scalar(k,
v,
global_step=step,
walltime=time.time())
# def summary_val(k, v):
# kwargs = {'tag': k, 'simple_value': float(v)}
# return self.tf.Summary.Value(**kwargs)
# summary = self.tf.Summary(value=[summary_val(k, v) for k, v in kvs.items()])
# event = self.event_pb2.Event(wall_time=time.time(), summary=summary)
# event.step = self.step # is there any reason why you'd want to specify the step?
# self.writer.WriteEvent(event)
# self.writer.Flush()
# self.step += 1
def close(self):
if self.writer:
self.writer.Close()
self.writer = None
