Repository: pfnet/pytorch-pfn-extras

Issue:
Do you have a plan to support TensorBoard
PyTorch supports TensorBoard naitively in `torch.utils.tensorboard`. But ppe's `LogReport` doesn't because tensorbord writers don't satisfy the requirement.

https://github.com/pfnet/pytorch-pfn-extras/blob/b2ddc4570f655bfe00f1d9f7dbb64fd4654a73c4/pytorch_pfn_extras/training/extensions/log_report.py#L57-L60

So my question is would ppe support it? Or do we have any other way to do it?


Discussion:
Let me look an easy way to work on this, maybe a custom extension?
Custom extension looks good to me. In my situation, this custom extension works well.

```
class TensorBoardReport(ppe.training.extensions.LogReport):
    def __call__(self, manager: ppe.training.ExtensionsManager):
        # accumulate the observations
        keys = self._keys
        observation = manager.observation
        updater = manager.updater
        summary = self._summary

        if keys is None:
            summary.add(observation)
        else:
            summary.add({k: observation[k] for k in keys if k in observation})

        assert self._writer is not None
        writer: torch.utils.tensorboad.SummaryWriter = self._writer

        if manager.is_before_training or self._trigger(manager):
            # output the result
            stats = self._summary.compute_mean()
            stats_cpu = {}
            for name, value in stats.items():
                stats_cpu[name] = float(value)  # copy to CPU

            stats_cpu["epoch"] = updater.epoch
            stats_cpu["iteration"] = updater.iteration
            stats_cpu["elapsed_time"] = manager.elapsed_time

            if self._postprocess is not None:
                self._postprocess(stats_cpu)

            self._log.append(stats_cpu)

            # write to the log file
            for key, value in stats_cpu.items():
                writer.add_scalar(key, value, stats_cpu["iteration"])

            # reset the summary for the next output
            self._init_summary()
```

But I don't know whether it supports all use cases.
I am thinking that we can expand our new writing system and create just a `TensorBoardWriter` object that takes care of this :)

Resolving patch:
diff --git a/pytorch_pfn_extras/writing.py b/pytorch_pfn_extras/writing.py
--- a/pytorch_pfn_extras/writing.py
+++ b/pytorch_pfn_extras/writing.py
@@ -500,3 +500,34 @@ def create_queue(self):
 
     def create_consumer(self, q):
         return multiprocessing.Process(target=self.consume, args=(q,))
+
+
+class TensorBoardWriter(object):
+    """ Writer that sends statistics to TensorBoard.
+
+    This class contains a `torch.utils.tensorboard.SummaryWriter`
+    object that is used to send the collected statistics to TensorBoard.
+    A list of stats can be specified to report only the desired ones.
+
+    """
+    def __init__(
+            self, savefun=None, fs=None, out_dir=None, stats=None, **kwds):
+        import torch.utils.tensorboard
+        self._stats = stats
+        self._writer = torch.utils.tensorboard.SummaryWriter(
+            log_dir=out_dir, **kwds)
+
+    def __call__(self, filename, out_dir, target, *, savefun=None):
+        stats_cpu = target
+        # we only take the last value
+        if isinstance(target, list):
+            stats_cpu = target[-1]
+
+        if not isinstance(stats_cpu, dict):
+            raise TypeError('target must be dict or list of dicts')
+        keys = stats_cpu.keys()
+        if self._stats is not None:
+            keys = self._stats
+        for key in keys:
+            value = stats_cpu[key]
+            self._writer.add_scalar(key, value, stats_cpu['iteration'])


Test patch:
diff --git a/.flexci/test.sh b/.flexci/test.sh
--- a/.flexci/test.sh
+++ b/.flexci/test.sh
@@ -4,7 +4,7 @@ perl -pi.bak -e 's|http://archive\.ubuntu\.com/ubuntu/|mirror://mirrors.ubuntu.c
 apt update
 apt -y install python3 python3-pip
 
-pip3 install torch torchvision pytorch-ignite pytest flake8 matplotlib onnx
+pip3 install torch torchvision pytorch-ignite pytest flake8 matplotlib tensorboard onnx
 pip3 install -e .
 
 # Run unit tests
diff --git a/tests/pytorch_pfn_extras_tests/test_writing.py b/tests/pytorch_pfn_extras_tests/test_writing.py
new file mode 100644
--- /dev/null
+++ b/tests/pytorch_pfn_extras_tests/test_writing.py
@@ -0,0 +1,20 @@
+import tempfile
+import os
+
+import pytest
+
+import pytorch_pfn_extras as ppe
+
+
+def test_tensorboard_writing():
+    pytest.importorskip('tensorboard')
+    data = {"a": 1, "iteration": 1}
+    with tempfile.TemporaryDirectory() as tempd:
+        writer = ppe.writing.TensorBoardWriter(
+            out_dir=tempd, filename_suffix='_test')
+        writer(None, None, data)
+        # Check that the file was generated
+        for snap in os.listdir(tempd):
+            assert '_test' in snap
+        writer._writer.flush()
+        writer._writer.close()