📡 📉 Adding Tensorboard Tracker

docs/source/tutorial/trackers/index.rst

@@ -6,4 +6,5 @@ Tracking Results during Training
    using_mlflow
    using_neptune
    using_wandb
+   using_tensorboard
    using_file

docs/source/tutorial/trackers/using_tensorboard.rst

@@ -0,0 +1,41 @@
+Using Tensorboard
+=========================
+
+`Tensorboard <https://www.tensorflow.org/tensorboard/>`_ (TB) is a service for tracking experimental results during training.
+It is part of the Tensorflow project.
+
+Minimal Pipeline Example
+---------------------------------
+A CSV log file can be generated with the following:
+
+.. code-block:: python
+
+    from pykeen.pipeline import pipeline
+
+    pipeline_result = pipeline(
+        model='RotatE',
+        dataset='Kinships',
+        result_tracker='tensorboard',
+    )
+
+It is placed in a subdirectory of :mod:`pystow` default data directory with PyKEEN called ``tensorboard``,
+which will likely be at ``~/.data/pykeen/logs/tensorboard`` on your system. The file is named based on the
+current time if no alternative is provided.
+
+Specifying a Name
+-----------------
+If you want to specify the name of the log file in the default directory, use the ``experiment_name`` keyword
+argument like:
+
+.. code-block:: python
+
+    from pykeen.pipeline import pipeline
+
+    pipeline_result = pipeline(
+        model='RotatE',
+        dataset='Kinships',
+        result_tracker='tensorboard',
+        result_tracker_kwargs=dict(
+        experiment_name='test',
+        ),
+    )

src/pykeen/trackers/__init__.py

@@ -7,6 +7,7 @@
 from .base import ResultTracker
 from .file import CSVResultTracker, FileResultTracker, JSONResultTracker
 from .mlflow import MLFlowResultTracker
+from .tensorboard import TensorBoardResultTracker
 from .neptune import NeptuneResultTracker
 from .wandb import WANDBResultTracker
 
@@ -20,6 +21,7 @@
     'WANDBResultTracker',
     'JSONResultTracker',
     'CSVResultTracker',
+    'TensorBoardResultTracker',
     # Utilities
     'tracker_resolver',
 ]

src/pykeen/trackers/tensorboard.py

@@ -0,0 +1,78 @@
+# -*- coding: utf-8 -*-
+
+"""An adapter for TensorBoard."""
+
+import datetime
+import pathlib
+from typing import Any, Dict, Mapping, Optional, TYPE_CHECKING, Type
+
+from .base import ResultTracker
+from ..constants import PYKEEN_LOGS
+from ..utils import flatten_dictionary
+
+if TYPE_CHECKING:
+    import torch.utils.tensorboard
+
+__all__ = [
+    'TensorBoardResultTracker',
+]
+
+
+class TensorBoardResultTracker(ResultTracker):
+    """A tracker for TensorBoard."""
+
+    #: The class that's used to instantiate a summarywriter
+    SummaryWriter: Type['torch.utils.tensorboard.SummaryWriter']
+
+    def __init__(
+        self,
+        experiment_path: Optional[str] = None,
+        experiment_name: Optional[str] = None,
+        tags: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize result tracking via Tensorboard.
+
+        :param experiment_path:
+            The experiment path. A custom path at which the tensorboard logs will be saved.
+        :param experiment_name:
+            The name of the experiment, will be used as a sub directory name for the logging. If no default is given,
+            the current time is used. If set, experiment_path is set, this argument has no effect.
+        :param tags:
+            The additional run details which are presented as tags to be logged
+        """
+        from torch.utils.tensorboard import SummaryWriter as _SummaryWriter
+        self.summary_writer_cls = _SummaryWriter
+        self.tags = tags
+
+        if experiment_path is None:
+            if experiment_name is None:
+                experiment_name = datetime.datetime.now().isoformat()
+            path = PYKEEN_LOGS.joinpath("tensorboard", experiment_name)
+        elif isinstance(experiment_path, str):
+            path = pathlib.Path(experiment_path)
+        self.path = path
+
+    def start_run(self, run_name: Optional[str] = None) -> None:  # noqa: D102
+        self.writer = self.SummaryWriter(log_dir=self.path, comment=run_name)
+
+    def log_metrics(
+        self,
+        metrics: Mapping[str, float],
+        step: Optional[int] = None,
+        prefix: Optional[str] = None,
+    ) -> None:  # noqa: D102
+        metrics = flatten_dictionary(dictionary=metrics, prefix=prefix)
+        for key, value in metrics.items():
+            self.writer.add_scalar(tag=key, scalar_value=value, global_step=step)
+        self.writer.flush()
+
+    def log_params(self, params: Mapping[str, Any], prefix: Optional[str] = None) -> None:  # noqa: D102
+        params = flatten_dictionary(dictionary=params, prefix=prefix)
+        for key, value in params.items():
+            self.writer.add_text(tag=str(key), text_string=str(value))
+        self.writer.flush()
+
+    def end_run(self) -> None:  # noqa: D102
+        self.writer.flush()
+        self.writer.close()