From f6d8928bd6849ce2a15a94ee18d01bebb3a6c20d Mon Sep 17 00:00:00 2001 From: evonneng Date: Tue, 31 May 2022 13:10:01 -0700 Subject: [PATCH] removed group and prefixes for logging (#118) --- radiance/nerf/graph/instant_ngp.py | 14 ++--- radiance/nerf/graph/nerfw.py | 2 +- radiance/nerf/graph/semantic_nerf.py | 2 +- radiance/nerf/graph/vanilla_nerf.py | 16 +++--- radiance/nerf/optimizers.py | 2 +- radiance/nerf/trainer.py | 2 +- radiance/utils/writer.py | 84 +++++++--------------------- 7 files changed, 40 insertions(+), 82 deletions(-) diff --git a/radiance/nerf/graph/instant_ngp.py b/radiance/nerf/graph/instant_ngp.py index d2dd45171..3c033a515 100644 --- a/radiance/nerf/graph/instant_ngp.py +++ b/radiance/nerf/graph/instant_ngp.py @@ -130,12 +130,12 @@ class NGPGraph(Graph): combined_acc = torch.cat([acc], dim=1) combined_depth = torch.cat([depth], dim=1) - writer.put_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img") - writer.put_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation") - writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth") + writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step) + writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step) + writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step) depth = visualization.apply_depth_colormap(outputs["depth_occupancy_grid"]) - writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth_occupancy_grid") + writer.put_image(name=f"depth_occupancy_grid/image_idx_{image_idx}", image=combined_depth, step=step) # Switch images from [H, W, C] to [1, C, H, W] for metrics computations image = torch.moveaxis(image, -1, 0)[None, ...] @@ -145,8 +145,8 @@ class NGPGraph(Graph): ssim = self.ssim(image, rgb) lpips = self.lpips(image, rgb) - writer.put_scalar(name=f"val_{image_idx}-fine", scalar=float(psnr), step=step, group="psnr") - writer.put_scalar(name=f"val_{image_idx}", scalar=float(ssim), step=step, group="ssim") - writer.put_scalar(name=f"val_{image_idx}", scalar=float(lpips), step=step, group="lpips") + writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(psnr), step=step) + writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(ssim), step=step) + writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(lpips), step=step) return psnr.item() diff --git a/radiance/nerf/graph/nerfw.py b/radiance/nerf/graph/nerfw.py index 67c603e10..42cf93335 100644 --- a/radiance/nerf/graph/nerfw.py +++ b/radiance/nerf/graph/nerfw.py @@ -210,4 +210,4 @@ class NerfWGraph(NeRFGraph): row2 = torch.cat([depth_fine, depth_fine_static, depth_coarse], dim=-2) combined_image = torch.cat([row0, row1, row2], dim=-3) - writer.put_image(name="image_idx_{image_idx}-nerfw", image=combined_image, step=step, group="img") + writer.put_image(name="img/image_idx_{image_idx}-nerfw", image=combined_image, step=step) diff --git a/radiance/nerf/graph/semantic_nerf.py b/radiance/nerf/graph/semantic_nerf.py index 46b326ce6..9ffb93136 100644 --- a/radiance/nerf/graph/semantic_nerf.py +++ b/radiance/nerf/graph/semantic_nerf.py @@ -155,4 +155,4 @@ class SemanticNerfGraph(NeRFGraph): self.stuff_classes ) semantic_labels_image = visualization.apply_colormap(semantic_labels[..., None]) - writer.put_image(name="image_idx_{image_idx}", image=semantic_labels_image, step=step, group="semantics") + writer.put_image(name="semantics/image_idx_{image_idx}", image=semantic_labels_image, step=step) diff --git a/radiance/nerf/graph/vanilla_nerf.py b/radiance/nerf/graph/vanilla_nerf.py index c7e5ba6eb..5e00795d9 100644 --- a/radiance/nerf/graph/vanilla_nerf.py +++ b/radiance/nerf/graph/vanilla_nerf.py @@ -152,10 +152,10 @@ class NeRFGraph(Graph): combined_depth = torch.cat([depth_coarse, depth_fine], dim=1) mask = visualization.apply_depth_colormap(mask[..., None]) - writer.put_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img") - writer.put_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation") - writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth") - writer.put_image(name=f"image_idx_{image_idx}", image=mask, step=step, group="mask") + writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step) + writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step) + writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step) + writer.put_image(name=f"mask/image_idx_{image_idx}", image=mask, step=step) # Switch images from [H, W, C] to [1, C, H, W] for metrics computations image = torch.moveaxis(image, -1, 0)[None, ...] @@ -167,10 +167,10 @@ class NeRFGraph(Graph): fine_ssim = self.ssim(image, rgb_fine) fine_lpips = self.lpips(image, rgb_fine) - writer.put_scalar(name=f"val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step, group="psnr") - writer.put_scalar(name=f"val_{image_idx}-fine", scalar=float(fine_psnr), step=step, group="psnr") - writer.put_scalar(name=f"val_{image_idx}", scalar=float(fine_ssim), step=step, group="ssim") - writer.put_scalar(name=f"val_{image_idx}", scalar=float(fine_lpips), step=step, group="lpips") + writer.put_scalar(name=f"psnr/val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step) + writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(fine_psnr), step=step) + writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(fine_ssim), step=step) + writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(fine_lpips), step=step) writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step) diff --git a/radiance/nerf/optimizers.py b/radiance/nerf/optimizers.py index 4606f0bd8..cddd02288 100644 --- a/radiance/nerf/optimizers.py +++ b/radiance/nerf/optimizers.py @@ -93,4 +93,4 @@ class Optimizers: scheduler.step() # TODO(ethan): clean this up. why is there indexing into a list? lr = scheduler.get_last_lr()[0] - writer.put_scalar(name=f"{param_group_name}", scalar=lr, step=step, group="learning_rate") + writer.put_scalar(name=f"learning_rate/{param_group_name}", scalar=lr, step=step) diff --git a/radiance/nerf/trainer.py b/radiance/nerf/trainer.py index e2f625203..2457f2458 100644 --- a/radiance/nerf/trainer.py +++ b/radiance/nerf/trainer.py @@ -230,7 +230,7 @@ class Trainer: ) if step != 0 and step % self.config.logging.steps_per_log == 0: - writer.put_dict(name="loss_dict", scalar_dict=loss_dict, step=step, group="Loss", prefix="train-") + writer.put_dict(name="Loss/train-loss_dict", scalar_dict=loss_dict, step=step) # TODO: add the learning rates to tensorboard/logging if step != 0 and self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0: self.save_checkpoint(self.config.graph.model_dir, step) diff --git a/radiance/utils/writer.py b/radiance/utils/writer.py index daf094534..22a1dfbb7 100644 --- a/radiance/utils/writer.py +++ b/radiance/utils/writer.py @@ -45,40 +45,27 @@ class EventType(enum.Enum): @check_main_thread -def put_image(name, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None): +def put_image(name, image: TensorType["H", "W", "C"], step: int): """Setter function to place images into the queue to be written out""" if isinstance(name, EventName): name = name.value - EVENT_STORAGE.append( - {"name": name, "write_type": EventType.IMAGE, "event": image, "step": step, "group": group, "prefix": prefix} - ) + EVENT_STORAGE.append({"name": name, "write_type": EventType.IMAGE, "event": image, "step": step}) @check_main_thread -def put_scalar(name: str, scalar: float, step: int, group: str = None, prefix: str = None): +def put_scalar(name: str, scalar: float, step: int): """Setter function to place scalars into the queue to be written out""" if isinstance(name, EventName): name = name.value - EVENT_STORAGE.append( - {"name": name, "write_type": EventType.SCALAR, "event": scalar, "step": step, "group": group, "prefix": prefix} - ) + EVENT_STORAGE.append({"name": name, "write_type": EventType.SCALAR, "event": scalar, "step": step}) @check_main_thread -def put_dict(name: str, scalar_dict: float, step: int, group: str = None, prefix: str = None): +def put_dict(name: str, scalar_dict: float, step: int): """Setter function to place a dictionary of scalars into the queue to be written out""" - EVENT_STORAGE.append( - { - "name": name, - "write_type": EventType.DICT, - "event": scalar_dict, - "step": step, - "group": group, - "prefix": prefix, - } - ) + EVENT_STORAGE.append({"name": name, "write_type": EventType.DICT, "event": scalar_dict, "step": step}) @check_main_thread @@ -87,8 +74,6 @@ def put_time( start_time: float, end_time: float, step: int, - group: str = None, - prefix: str = None, avg_over_iters: bool = False, avg_over_batch: int = None, update_eta: bool = False, @@ -116,16 +101,16 @@ def put_time( curr_buffer.pop(0) curr_buffer.append(val) curr_avg = sum(curr_buffer) / len(curr_buffer) - put_scalar(name, curr_avg, step, group, prefix) + put_scalar(name, curr_avg, step) GLOBAL_BUFFER["events"][name] = {"buffer": curr_buffer, "avg": curr_avg} else: - put_scalar(name, val, step, group, prefix) + put_scalar(name, val, step) if update_eta: ## NOTE: eta should be called with avg train iteration time remain_iter = GLOBAL_BUFFER["max_iter"] - step remain_time = remain_iter * GLOBAL_BUFFER["events"][name]["avg"] - put_scalar(EventName.ETA, remain_time, step, group, prefix) + put_scalar(EventName.ETA, remain_time, step) @check_main_thread @@ -135,11 +120,11 @@ def write_out_storage(): for event in EVENT_STORAGE: write_func = getattr(writer, event["write_type"].value) if event["write_type"] == EventType.DICT: - write_func(event["event"], event["step"], event["group"], event["prefix"]) + write_func(event["event"], event["step"]) if isinstance(writer, LocalWriter): continue else: - write_func(event["name"], event["event"], event["step"], event["group"], event["prefix"]) + write_func(event["name"], event["event"], event["step"]) if isinstance(writer, LocalWriter): break EVENT_STORAGE.clear() @@ -165,16 +150,6 @@ def setup_event_writers(config: DictConfig) -> None: GLOBAL_BUFFER["events"] = {} -def get_tensorboard_name(name: str, group: str = None, prefix: str = None): - """Returns a string for tensorboard with an optional group and prefix. - Where tensorboard_name has the form `group/prefix-name`. - """ - group_string = f"{group}/" if group else "" - prefix_string = f"{prefix}" if prefix else "" - tensorboard_name = f"{group_string}{prefix_string}{name}" - return tensorboard_name - - class Writer: """Writer class""" @@ -182,46 +157,36 @@ class Writer: self.save_dir = save_dir @abstractmethod - def write_image( - self, name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None - ) -> None: + def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: """_summary_ Args: name (str): data identifier image (TensorType["H", "W", "C"]): 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_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None: + def write_scalar(self, name: str, scalar: float, step: int) -> None: """Required method to write a single scalar value to the logger Args: name (str): data identifier 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 @check_main_thread - def write_scalar_dict( - self, scalar_dict: Dict[str, float], step: int, group: str = None, prefix: str = None - ) -> None: + def write_scalar_dict(self, scalar_dict: Dict[str, float], step: int) -> None: """Function that writes out all scalars from a given dictionary to the logger Args: 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-" """ for name, scalar in scalar_dict.items(): - self.write_scalar(name, scalar, step, group=group, prefix=prefix) + self.write_scalar(name, scalar, step) @decorate_all([check_main_thread]) @@ -232,9 +197,7 @@ class TensorboardWriter(Writer): super().__init__(save_dir) self.tb_writer = SummaryWriter(log_dir=self.save_dir) - def write_image( - self, name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None - ) -> None: + def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: """_summary_ Args: @@ -242,20 +205,17 @@ class TensorboardWriter(Writer): image (TensorType["H", "W", "C"]): rendered image to write """ image = to8b(image) - tensorboard_name = get_tensorboard_name(name, group=group, prefix=prefix) - self.tb_writer.add_image(tensorboard_name, image, step, dataformats="HWC") + self.tb_writer.add_image(name, image, step, dataformats="HWC") - def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None: + def write_scalar(self, name: str, scalar: float, step: int) -> None: """Tensorboard 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 - group (str)): a prefix to group tensorboard scalars """ - tensorboard_name = get_tensorboard_name(name, group=group, prefix=prefix) - self.tb_writer.add_scalar(tensorboard_name, scalar, step) + self.tb_writer.add_scalar(name, scalar, step) def _cursorup(x: int): @@ -303,15 +263,13 @@ class LocalWriter(Writer): self.keys = set() self.past_mssgs = ["", ""] - def write_image( - self, name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None - ) -> None: + def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: if name in self.stats_to_track: image = to8b(image) image_path = os.path.join(self.save_dir, f"{name}.jpg") imageio.imwrite(image_path, np.uint8(image.cpu().numpy() * 255.0)) - def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None: + def write_scalar(self, name: str, scalar: float, step: int) -> None: if step > 0: latest_map, new_key = self._consolidate_events() self._update_header(latest_map, new_key)