diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index d1462303c..f022dd4df 100644 --- a/mattport/utils/writer.py +++ b/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)