def add_animated_gif( writer: SummaryWriter, tag: str, image_tensor: Union[np.ndarray, torch.Tensor], max_out: int = 3, frame_dim: int = -3, scale_factor: float = 1.0, global_step: Optional[int] = None, ) -> None: """Creates an animated gif out of an image tensor in 'CHWD' format and writes it with SummaryWriter. Args: writer: Tensorboard SummaryWriter to write to tag: Data identifier image_tensor: tensor for the image to add, expected to be in `CHWD` format max_out: maximum number of image channels to animate through frame_dim: the dimension used as frames for GIF image, expect input data shape as `CHWD`, default to `-3` (the first spatial dim) scale_factor: amount to multiply values by. If the image data is between 0 and 1, using 255 for this value will scale it to displayable range global_step: Global step value to record """ summary = make_animated_gif_summary(tag=tag, image=image_tensor, writer=writer, max_out=max_out, frame_dim=frame_dim, scale_factor=scale_factor) for s in summary: # add GIF for every channel separately writer._get_file_writer().add_summary(s, global_step)
def add_animated_gif_no_channels( writer: SummaryWriter, tag: str, image_tensor: Union[np.ndarray, torch.Tensor], max_out: int, scale_factor: float, global_step: Optional[int] = None, ) -> None: """Creates an animated gif out of an image tensor in 'HWD' format that does not have a channel dimension and writes it with SummaryWriter. This is similar to the "add_animated_gif" after inserting a channel dimension of 1. Args: writer: Tensorboard SummaryWriter to write to tag: Data identifier image_tensor: tensor for the image to add, expected to be in CHWD format max_out: maximum number of slices to animate through scale_factor: amount to multiply values by. If the image data is between 0 and 1, using 255 for this value will scale it to displayable range global_step: Global step value to record """ writer._get_file_writer().add_summary( make_animated_gif_summary(tag, image_tensor, max_out=max_out, animation_axes=[1], image_axes=[1, 2], scale_factor=scale_factor), global_step, )
def add_animated_gif( writer: SummaryWriter, tag: str, image_tensor: Union[np.array, Tensor], max_out: int, scale_factor: float, global_step: int = None, ): """Creates an animated gif out of an image tensor in 'CHWD' format and writes it with SummaryWriter. Args: writer (SummaryWriter): Tensorboard SummaryWriter to write to tag (str): Data identifier image_tensor (np.array or Tensor): tensor for the image to add, expected to be in CHWD format max_out (int): maximum number of slices to animate through scale_factor (float): amount to multiply values by. If the image data is between 0 and 1, using 255 for this value will scale it to displayable range global_step (int): Global step value to record """ writer._get_file_writer().add_summary( make_animated_gif_summary( tag, image_tensor, max_out=max_out, animation_axes=[1], image_axes=[2, 3], scale_factor=scale_factor ), global_step, )
def async_log(tensorboard_dir, log): global log_queue if log_queue is None or log is None: return finished = False tensorboard_writer = SummaryWriter(tensorboard_dir) while True: if finished and log_queue.empty(): tensorboard_writer.close() return temp = log_queue.get() if temp is None: return if temp == "finish": finished = True continue if temp[0] == "tensorboard": function, args = temp[1:] getattr(tensorboard_writer, function)(*args) tensorboard_writer._get_file_writer().flush() elif temp[0] == "log": log.info(temp[1])
class TensorboardWriter: def __init__(self, outdir): assert (os.path.isdir(outdir)) self.outdir = outdir self.writer = SummaryWriter(self.outdir, flush_secs=10) def save_scalar(self, tag, scalar_value, global_step=None): self.writer.add_scalar(tag, scalar_value, global_step) def save_scalars(self, main_tag, scalars_dict, global_step=None): self.writer.add_scalars(main_tag, scalars_dict, global_step) def save_image(self, tag, image, global_step=None, dataformats='CHW'): self.writer.add_image(tag=tag, img_tensor=image, global_step=global_step, dataformats=dataformats) def save_figure(self, tag, figure, global_step=None, close=False): self.writer.add_figure(tag=tag, figure=figure, global_step=global_step, close=close) def save_graph(self, model, inputs_to_model=None): """ Saves graph to the tensorboard. Ideally call once. :param model: The torch.nn.Module object :param inputs_to_model: tensor or a list of tensor(batch) will also be showed. :return: None """ try: self.writer.add_graph(model, inputs_to_model) except Exception as e: logger.exception( 'Check this for fix: https://github.com/lanpa/tensorboardX/issues/389#issuecomment-475879228' ) def save_embedding(self, mat, metadata=None, label_img=None, global_step=None, tag='default', metadata_header=None): self.writer.add_embedding(mat, metadata, label_img, global_step, tag, metadata_header) def flush(self): """ If you need to flush all data immediately. """ self.writer._get_file_writer().flush() def close(self): """ To be called in the end """ self.writer.close()
class NestedFolderTensorboardLogger(LightningLoggerBase): def __init__(self, save_dir, name, **kwargs): super().__init__() self.save_dir = save_dir self._name = name self.experiment_root = None self.experiment_discriminator = None self.experiment_generator = None self.kwargs = kwargs self.setup_experiments() def setup_experiments(self): root_dir = os.path.join(self.save_dir, self.name) os.makedirs(root_dir, exist_ok=True) generator_dir = os.path.join(self.save_dir, self.name, "generator") discriminator_dir = os.path.join(self.save_dir, self.name, "discriminator") os.makedirs(generator_dir, exist_ok=True) os.makedirs(discriminator_dir, exist_ok=True) self.experiment_root = SummaryWriter(log_dir=root_dir, **self.kwargs) self.experiment_discriminator = SummaryWriter( log_dir=discriminator_dir, **self.kwargs) self.experiment_generator = SummaryWriter(log_dir=generator_dir, **self.kwargs) @rank_zero_only def log_hyperparams(self, params): params = vars(params) self.experiment_root.add_hparams(hparam_dict=dict(params), metric_dict={}) @rank_zero_only def log_metrics(self, metrics, step_num): return @rank_zero_only def log_image(self, step, visualization): self.experiment_root.add_image('visualization', visualization, step, dataformats='HWC') @rank_zero_only def log_generator_losses(self, step, gan_loss, reconstruction_loss): self.experiment_generator.add_scalars( 'G/losses', { 'gan_loss': gan_loss, 'reconstruction_loss': reconstruction_loss }, step) @rank_zero_only def log_discriminator_losses(self, step, real_loss, fake_loss): self.experiment_discriminator.add_scalars('D/losses', { 'real': real_loss, 'fake': fake_loss }, step) @rank_zero_only def log_total_generator_loss(self, step, loss): self.experiment_generator.add_scalar('total', loss, step) @rank_zero_only def log_total_discriminator_loss(self, step, loss): self.experiment_discriminator.add_scalar('total', loss, step) @rank_zero_only def save(self): self.experiment_generator._get_file_writer().flush() self.experiment_discriminator._get_file_writer().flush() self.experiment_root._get_file_writer().flush() @rank_zero_only def finalize(self, status): self.save() @property def version(self): return self._name @property def name(self): return self._name
class NestedTensorboardLogger(LightningLoggerBase): @property def experiment(self): return self.experiment_root @property def save_dir(self): return self._save_dir def __init__(self, save_dir, name, **kwargs): super().__init__() self._save_dir = save_dir self._name = name self.experiment_root = None self.experiment_train = None self.experiment_val = None self.kwargs = kwargs self.setup_experiments() def setup_experiments(self): root_dir = os.path.join(self.save_dir, self.name) os.makedirs(root_dir, exist_ok=True) train_dir = os.path.join(self.save_dir, self.name, "train") val_dir = os.path.join(self.save_dir, self.name, "val") os.makedirs(train_dir, exist_ok=True) os.makedirs(val_dir, exist_ok=True) self.experiment_root = SummaryWriter(log_dir=root_dir, **self.kwargs) self.experiment_val = SummaryWriter(log_dir=val_dir, **self.kwargs) self.experiment_train = SummaryWriter(log_dir=train_dir, **self.kwargs) def log_hyperparams(self, params): for k in params: if params[k] is None: params[k] = "None" if type(params[k]) == list: params[k] = torch.LongTensor(params[k]) self.experiment_root.add_hparams(hparam_dict=dict(params), metric_dict={}) def log_train_loss(self, global_step, error): self.experiment_train.add_scalar('loss', error, global_step) def log_val_loss(self, global_step, loss): self.experiment_val.add_scalar('loss', loss, global_step) def log_metrics(self, metrics, step): return def save(self): self.experiment_train._get_file_writer().flush() self.experiment_val._get_file_writer().flush() self.experiment_root._get_file_writer().flush() def finalize(self, status): self.save() @property def version(self): return self._name @property def name(self): return self._name