fixed writer

mattport/utils/writer.py

@@ -26,41 +26,56 @@ class Writer:
         self.is_main_thread = local_rank % world_size == 0
 
     @abstractmethod
-    def write_image(self, name: str, x: TensorType["H", "W", 3], step: int) -> None:
+    def write_image(
+        self, name: str, x: TensorType["H", "W", 3], step: int, group: str = None, prefix: str = None
+    ) -> None:
         """_summary_
 
         Args:
             name (str): data identifier
             x (TensorType["H", "W", 3]): rendered image to write
+            step (int): the time step to log
+            group (str): the group e.g., "Loss", "Accuracy", "Time"
+            prefix (str): the prefix e.g., "train-", "test-"
         """
         raise NotImplementedError
 
-    @abstractmethod
-    def write_text(self, name: str, x: str) -> None:
+    @check_main_thread
+    def write_text(self, name: str, x: str, step: int, group: str = None, prefix: str = None) -> None:
         """Required method to write a string to summary
 
         Args:
             name (str): data identifier
             x (str): string to write
+            step (int): the time step to log
+            group (str): the group e.g., "Loss", "Accuracy", "Time"
+            prefix (str): the prefix e.g., "train-", "test-"
         """
         raise NotImplementedError
 
     @abstractmethod
-    def write_scalar(self, name: str, scalar: float, step: float, group: str = None) -> None:
+    def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None:
         """Required method to write a single scalar value to the logger
 
         Args:
             name (str): data identifier
-            x (float): x value to write
-            y (float): y value to write
+            step (int): the time step to log
+            group (str): the group e.g., "Loss", "Accuracy", "Time"
+            prefix (str): the prefix e.g., "train-", "test-"
         """
         raise NotImplementedError
 
-    def write_scalar_dict(self, scalar_dict: Dict[str, float], step: int, group: str = None,  prefix: str = None) -> None:
+    @check_main_thread
+    def write_scalar_dict(
+        self, scalar_dict: Dict[str, float], step: int, group: str = None, prefix: str = None
+    ) -> None:
         """Function that writes out all scalars from a given dictionary to the logger
 
         Args:
-            scalar_dict (dict): dictionary containing all scalar values
+            scalar_dict (dict): dictionary containing all scalar values with key names and quantities
+            step (int): the time step to log
+            group (str): the group e.g., "Loss", "Accuracy", "Time"
+            prefix (str): the prefix e.g., "train-", "test-"
         """
         if self.is_main_thread:
             for name, scalar in scalar_dict.items():
@@ -76,7 +91,9 @@ class TensorboardWriter(Writer):
             self.tb_writer = SummaryWriter(log_dir=self.save_dir)
 
     @check_main_thread
-    def write_image(self, name: str, x: TensorType["H", "W", 3], step: int) -> None:
+    def write_image(
+        self, name: str, x: TensorType["H", "W", 3], step: int, group: str = None, prefix: str = None
+    ) -> None:
         """_summary_
 
         Args:
@@ -87,17 +104,7 @@ class TensorboardWriter(Writer):
         self.tb_writer.add_image(name, x, step, dataformats="HWC")
 
     @check_main_thread
-    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
-        """
-        self.tb_writer.add_text(name, x)
-
-    @check_main_thread
-    def write_scalar(self, name: str, scalar: float, step: float, group: str = None, prefix: str = None) -> None:
+    def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None:
         """Tensorboard method to write a single scalar value to the logger
 
         Args:
@@ -115,33 +122,22 @@ class LocalWriter(Writer):
     """Local Writer Class"""
 
     @check_main_thread
-    def write_image(self, name: str, x: TensorType["H", "W", 3]) -> None:
-        """Logs image to a jpg file in save_directory
-
-        Args:
-            name (str): data identifier to be used as file name
-            x (TensorType["H", "W", 3]): rendered image to write
-        """
+    def write_image(
+        self, name: str, x: TensorType["H", "W", 3], step: int, group: str = None, prefix: str = None
+    ) -> None:
         x = to8b(x)
         image_path = os.path.join(self.save_dir, f"{name}.jpg")
         imageio.imwrite(image_path, np.uint8(x.cpu().numpy() * 255.0))
 
     @check_main_thread
-    def write_text(self, name: str, x: str) -> None:
+    def write_text(self, name: str, x: str, step: int, group: str = None, prefix: str = None) -> None:
         """Logs text locally to the terminal in place
         Args:
             name (str): data identifier
             x (str): string to write
         """
-        print(f"{name}: {x}", end="\r")
+        print(f"{name}: {x}, step: {step}", end="\r")
 
     @check_main_thread
-    def write_scalar(self, name: str, x: float, y: float) -> None:
-        """Logs scalar locally to the terminal in place
-
-        Args:
-            name (str): data identifier
-            x (float): x value to write
-            y (float): y value to write
-        """
-        self.write_text(name, f"x: {x:04f} y: {y:04f}")
+    def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None:
+        self.write_text(name, f"scalar: {scalar:04f}", step)