class TensorboardWriter: def __init__(self, log_dir: str, *args: Any, **kwargs: Any): r"""A Wrapper for tensorboard SummaryWriter. It creates a dummy writer when log_dir is empty string or None. It also has functionality that generates tb video directly from numpy images. Args: log_dir: Save directory location. Will not write to disk if log_dir is an empty string. *args: Additional positional args for SummaryWriter **kwargs: Additional keyword args for SummaryWriter """ self.writer = None if log_dir is not None and len(log_dir) > 0: self.writer = SummaryWriter(log_dir, *args, **kwargs) def __getattr__(self, item): if self.writer: return self.writer.__getattribute__(item) else: return lambda *args, **kwargs: None def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): if self.writer: self.writer.close()
class TensorboardWriter: def __init__(self, log_dir: str, *args: Any, **kwargs: Any): r"""A Wrapper for tensorboard SummaryWriter. It creates a dummy writer when log_dir is empty string or None. It also has functionality that generates tb video directly from numpy images. Args: log_dir: Save directory location. Will not write to disk if log_dir is an empty string. *args: Additional positional args for SummaryWriter **kwargs: Additional keyword args for SummaryWriter """ self.writer = None if log_dir is not None and len(log_dir) > 0: self.writer = SummaryWriter(log_dir, *args, **kwargs) def __getattr__(self, item): if self.writer: return self.writer.__getattribute__(item) else: return lambda *args, **kwargs: None def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): if self.writer: self.writer.close() def add_video_from_np_images(self, video_name: str, step_idx: int, images: np.ndarray, fps: int = 10) -> None: r"""Write video into tensorboard from images frames. Args: video_name: name of video string. step_idx: int of checkpoint index to be displayed. images: list of n frames. Each frame is a np.ndarray of shape. fps: frame per second for output video. Returns: None. """ if not self.writer: return # initial shape of np.ndarray list: N * (H, W, 3) frame_tensors = [ torch.from_numpy(np_arr).unsqueeze(0) for np_arr in images ] video_tensor = torch.cat(tuple(frame_tensors)) video_tensor = video_tensor.permute(0, 3, 1, 2).unsqueeze(0) # final shape of video tensor: (1, n, 3, H, W) self.writer.add_video(video_name, video_tensor, fps=fps, global_step=step_idx)
class TensorboardLogger(object): def __init__(self): self._logger = None self._global_step = 0 def create(self, path): self._logger = SummaryWriter(path) def noop(self, *args, **kwargs): return def step(self): self._global_step += 1 @property def global_step(self): return self._global_step def log_scaler_dict(self, log_dict, prefix=''): """ log a dictionary of scalar values""" if self._logger is None: return if prefix: prefix = f'{prefix}_' for name, value in log_dict.items(): if isinstance(value, dict): self.log_scaler_dict(value, self._global_step, prefix=f'{prefix}{name}') else: self._logger.add_scalar(f'{prefix}{name}', value, self._global_step) def __getattr__(self, name): if self._logger is None: return self.noop return self._logger.__getattribute__(name)
class Logger(object): def __init__(self, log_dir=None, tensorboard=False, txt=False, logfile=None, **kwargs): self.tblog = None if tensorboard: try: self.tblog = SummaryWriter(log_dir=log_dir, **kwargs) except: self.tblog = None self.txtlog = None if txt: if logfile is None: logfile = open(os.path.join(log_dir, 'logfile.txt'), 'a+') self._old_stdout = sys.stdout self.txtlog = Tee(self._old_stdout, logfile) sys.stdout = self.txtlog now = time.strftime("%b-%d-%Y-%H%M%S") title = '**** Beginning Log {} ****\n'.format(now) title_stars = '*' * (len(title) - 1) + '\n' self.txtlog.write(title_stars + title + title_stars, logonly=True) self.global_step = None self.tag_fmt = None def set_step(self, step): self.global_step = step def get_step(self): return self.global_step def set_tag_format(self, fmt=None): self.tag_fmt = fmt def get_tag_format(self): return self.tag_fmt def add_hparams(self, param_dict, metrics={}): self.tblog.add_hparams(param_dict, metrics) def add(self, data_type, tag, *args, global_step=None, **kwargs): # TODO: test and maybe clean if self.tblog is None: return None # if data_type == 'scalar' and not isinstance(args[0], float): # print(tag, type(args[0]), args[0]) # assert False add_fn = self.tblog.__getattribute__('add_{}'.format(data_type)) if global_step is None: global_step = self.global_step if self.tag_fmt is not None: tag = self.tag_fmt.format(tag) add_fn(tag, *args, global_step=global_step, **kwargs) def flush(self): if self.tblog is not None: self.tblog.flush() if self.txtlog is not None: self.txtlog.flush()