From 89a7d1ba531f946dabb8ed96ddd0a42d258059ef Mon Sep 17 00:00:00 2001 From: evonneng Date: Tue, 19 Apr 2022 07:02:31 +0000 Subject: [PATCH 1/2] 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) From 5b629b94956aa63408c68444db98977319ad7fb6 Mon Sep 17 00:00:00 2001 From: evonneng Date: Wed, 20 Apr 2022 04:45:49 +0000 Subject: [PATCH 2/2] added writer with declarator to checkif you are on main thread; add write image --- mattport/utils/writer.py | 48 ++++++++++++++++++++++++++++++++++------ 1 file changed, 41 insertions(+), 7 deletions(-) diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index d2c95a3d0..76405f4c5 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -4,8 +4,21 @@ Generic Writer class from abc import abstractmethod +from typing import Callable +import torch from torch.utils.tensorboard import SummaryWriter +from torchtyping import TensorType + + +def require_main_thread(func: Callable) -> Callable: + """Declarator: check if you are on main thread""" + + def wrapper(self, *args, **kwargs): + if self.is_main_thread: + func(self, *args, **kwargs) + + return wrapper class Writer: @@ -16,6 +29,16 @@ class Writer: self.writer = None self.is_main_thread = local_rank % world_size == 0 + @abstractmethod + def write_image(self, name: str, x: TensorType["H", "W", 3]) -> None: + """_summary_ + + Args: + name (str): data identifier + x (TensorType["H", "W", 3]): rendered image to write + """ + raise NotImplementedError + @abstractmethod def write_text(self, name: str, x: str) -> None: """Required method to write a string to summary @@ -31,7 +54,7 @@ class Writer: """Required method to write a single scalar value to the logger Args: - name (str): name of chart to which you log (x,y) + name (str): data identifier x (float): x value to write y (float): y value to write """ @@ -56,6 +79,19 @@ class TensorboardWriter(Writer): if self.is_main_thread: self.writer = SummaryWriter(log_dir=self.save_dir) + @require_main_thread + def write_image(self, name: str, x: TensorType["H", "W", 3]) -> None: + """_summary_ + + Args: + name (str): data identifier + x (TensorType["H", "W", 3]): rendered image to write + """ + to8b = lambda x: (255 * torch.clip(x, 0, 1)).to(torch.uint8) + x = to8b(x) + self.writer.add_images(name, x) + + @require_main_thread def write_text(self, name: str, x: str) -> None: """Tensorboard method to write a string to summary @@ -63,17 +99,15 @@ class TensorboardWriter(Writer): name (str): data identifier x (str): string to write """ - if self.is_main_thread: - self.writer.add_text(name, x) + self.writer.add_text(name, x) - @abstractmethod + @require_main_thread 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) + name (str): data identifier x (float): x value to write y (float): y value to write """ - if self.is_main_thread: - self.writer.add_scalar(name, x, y) + self.writer.add_scalar(name, x, y)