Logging via protobuf, without tensorflow

tensorboard_logger/tensorboard_logger.py

@@ -1,9 +1,11 @@
 # -*- coding: utf-8 -*-
 from collections import defaultdict
+import os
 import re
+import time
 
 import six
-import tensorflow as tf
+from .tf_protobuf import summary_pb2, event_pb2
 
 
 __all__ = ['Logger', 'configure', 'log_value']
@@ -15,16 +17,23 @@
 
 class Logger(object):
     def __init__(self, logdir, flush_secs=2, is_dummy=False):
-        self._session = tf.Session()
-        if not is_dummy:
-            self._writer = tf.train.SummaryWriter(logdir, flush_secs=flush_secs)
-        self._loggers = {}
+        self._name_to_tf_name = {}
         self._tf_names = set()
         self.is_dummy = is_dummy
+        self.logdir = logdir
+        self.flush_secs = flush_secs  # TODO
+        self._writer = None
         if is_dummy:
             self.dummy_log = defaultdict(list)
-
-    def log_value(self, name, value, step):
+        else:
+            if not os.path.exists(self.logdir):
+                os.makedirs(self.logdir)
+            filename = os.path.join(
+                self.logdir,
+                'events.out.tfevents.{}.ws'.format(int(time.time())))
+            self._writer = open(filename, 'wb')
+
+    def log_value(self, name, value, step=None):
         """ Log new value for given name on given step.
         value should be a real number (it will be converted to float),
         and name should be a string (it will be converted to a valid
@@ -42,16 +51,17 @@ def log_value(self, name, value, step):
                             .format(type(value)))
         value = float(value)
 
-        if not isinstance(step, six.integer_types):
+        if step is not None and not isinstance(step, six.integer_types):
             raise TypeError('"step" should be an integer, got {}'
                             .format(type(step)))
 
         try:
-            logger = self._loggers[name]
+            tf_name = self._name_to_tf_name[name]
         except KeyError:
             tf_name = self._make_tf_name(name)
-            logger = self._loggers[name] = self._make_logger(tf_name, value)
-        logger(value, step)
+            self._name_to_tf_name[name] = tf_name
+
+        self._log_value(tf_name, value, step)
 
     def _make_tf_name(self, name):
         tf_base_name = tf_name = make_valid_tf_name(name)
@@ -62,24 +72,25 @@ def _make_tf_name(self, name):
         self._tf_names.add(tf_name)
         return tf_name
 
-    def _make_logger(self, tf_name, value):
-        dtype = tf.float32
-        variable = tf.Variable(
-            initial_value=value, dtype=dtype, trainable=False, name=tf_name)
-        self._session.run(tf.initialize_variables([variable], tf_name))
-        summary_op = tf.scalar_summary(tf_name, variable)
-        new_value = tf.placeholder(dtype, shape=[])
-        assign_op = tf.assign(variable, new_value)
-
-        def logger(x, i):
-            _, summary = self._session.run([assign_op, summary_op],
-                                           {new_value: x})
-            if self.is_dummy:
-                self.dummy_log[tf_name].append((i, x))
-            else:
-                self._writer.add_summary(summary, i)
-
-        return logger
+    def _log_value(self, tf_name, value, step=None):
+        summary = summary_pb2.Summary(value=[
+            summary_pb2.Summary.Value(node_name=tf_name, simple_value=value)
+            ])
+        event = event_pb2.Event(wall_time=time.time(), summary=summary)
+        if step is not None:
+            event.step = int(step)
+        if self.is_dummy:
+            self.dummy_log[tf_name].append((step, value))
+        else:
+            self._write_event(event)
+
+    def _write_event(self, event):
+        self._writer.write(event.SerializeToString())
+        self._writer.flush()  # FIXME
+
+    def __del__(self):
+        if self._writer is not None:
+            self._writer.close()
 
 
 def make_valid_tf_name(name):