From 89a7d1ba531f946dabb8ed96ddd0a42d258059ef Mon Sep 17 00:00:00 2001
From: evonneng <evonne_ng@berkeley.edu>
Date: Tue, 19 Apr 2022 07:02:31 +0000
Subject: [PATCH] added base writer class and tensorboard support; todo: init
 from config in trainer.py

---
 mattport/utils/writer.py | 79 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)
 create mode 100644 mattport/utils/writer.py

diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py
new file mode 100644
index 000000000..d2c95a3d0
--- /dev/null
+++ b/mattport/utils/writer.py
@@ -0,0 +1,79 @@
+"""
+Generic Writer class
+"""
+
+
+from abc import abstractmethod
+
+from torch.utils.tensorboard import SummaryWriter
+
+
+class Writer:
+    """Writer class"""
+
+    def __init__(self, local_rank: int, world_size: int, save_dir: str):
+        self.save_dir = save_dir
+        self.writer = None
+        self.is_main_thread = local_rank % world_size == 0
+
+    @abstractmethod
+    def write_text(self, name: str, x: str) -> None:
+        """Required method to write a string to summary
+
+        Args:
+            name (str): data identifier
+            x (str): string to write
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def write_scalar(self, name: str, x: float, y: float) -> None:
+        """Required method to write a single scalar value to the logger
+
+        Args:
+            name (str): name of chart to which you log (x,y)
+            x (float): x value to write
+            y (float): y value to write
+        """
+        raise NotImplementedError
+
+    def write_scalar_dict(self, scalar_dict: dict) -> None:
+        """Function that writes out all scalars from a given dictionary to the logger
+
+        Args:
+            scalar_dict (dict): dictionary containing all scalar values
+        """
+        if self.is_main_thread:
+            for name, (x, y) in scalar_dict.items():
+                self.writer.write_scalar(name, x, y)
+
+
+class TensorboardWriter(Writer):
+    """Tensorboard Writer Class"""
+
+    def __init__(self, local_rank: int, world_size: int, save_dir: str):
+        super().__init__(local_rank, world_size, save_dir)
+        if self.is_main_thread:
+            self.writer = SummaryWriter(log_dir=self.save_dir)
+
+    def write_text(self, name: str, x: str) -> None:
+        """Tensorboard method to write a string to summary
+
+        Args:
+            name (str): data identifier
+            x (str): string to write
+        """
+        if self.is_main_thread:
+            self.writer.add_text(name, x)
+
+    @abstractmethod
+    def write_scalar(self, name: str, x: float, y: float) -> None:
+        """Tensorboard method to write a single scalar value to the logger
+
+        Args:
+            name (str): name of chart to which you log (x,y)
+            x (float): x value to write
+            y (float): y value to write
+        """
+        if self.is_main_thread:
+            self.writer.add_scalar(name, x, y)