Allow using a custom field in the label as the step when logging to tensorboard.
"""Utilities for logging to the terminal."""
import time
+from typing import Optional
+from absl import logging
from acme.utils.loggers import base
import tensorflow as tf
@@ -37,23 +39,34 @@ class TFSummaryLogger(base.Logger):
self,
logdir: str,
label: str = 'Logs',
+ steps_key: Optional[str] = None
):
"""Initializes the logger.
Args:
logdir: directory to which we should log files.
label: label string to use when logging. Default to 'Logs'.
+ steps_key: key to use for steps. Must be in the values passed to write.
"""
self._time = time.time()
self.label = label
self._iter = 0
self.summary = tf.summary.create_file_writer(logdir)
+ self._steps_key = steps_key
def write(self, values: base.LoggingData):
+ if self._steps_key is not None and self._steps_key not in values:
+ logging.warn('steps key %s not found. Skip logging.', self._steps_key)
+ return
+
+ step = values[
+ self._steps_key] if self._steps_key is not None else self._iter
+
with self.summary.as_default():
- for key, value in values.items():
+ # TODO(b/159065169): Remove this suppression once the bug is resolved.
+ # pytype: disable=unsupported-operands
+ for key in values.keys() - [self._steps_key]:
+ # pytype: enable=unsupported-operands
tf.summary.scalar(
- f'{self.label}/{_format_key(key)}',
- value,
- step=self._iter)
+ f'{self.label}/{_format_key(key)}', data=values[key], step=step)
self._iter += 1
