| | """
|
| | Dumps things to tensorboard and console
|
| | """
|
| |
|
| | import datetime
|
| | import logging
|
| | import math
|
| | import os
|
| | from collections import defaultdict
|
| | from pathlib import Path
|
| | from typing import Optional, Union
|
| | import matplotlib
|
| | matplotlib.use('TkAgg')
|
| | import matplotlib.pyplot as plt
|
| | import numpy as np
|
| | import torch
|
| | import torchaudio
|
| | from PIL import Image
|
| | from pytz import timezone
|
| | from torch.utils.tensorboard import SummaryWriter
|
| |
|
| | from .email_utils import EmailSender
|
| | from .time_estimator import PartialTimeEstimator, TimeEstimator
|
| | from .timezone import my_timezone
|
| |
|
| |
|
| | def tensor_to_numpy(image: torch.Tensor):
|
| | image_np = (image.numpy() * 255).astype('uint8')
|
| | return image_np
|
| |
|
| |
|
| | def detach_to_cpu(x: torch.Tensor):
|
| | return x.detach().cpu()
|
| |
|
| |
|
| | def fix_width_trunc(x: float):
|
| | return ('{:.9s}'.format('{:0.9f}'.format(x)))
|
| |
|
| |
|
| | def plot_spectrogram(spectrogram: np.ndarray, title=None, ylabel="freq_bin", ax=None):
|
| | if ax is None:
|
| | _, ax = plt.subplots(1, 1)
|
| | if title is not None:
|
| | ax.set_title(title)
|
| | ax.set_ylabel(ylabel)
|
| | ax.imshow(spectrogram, origin="lower", aspect="auto", interpolation="nearest")
|
| |
|
| |
|
| | class TensorboardLogger:
|
| |
|
| | def __init__(self,
|
| | exp_id: str,
|
| | run_dir: Union[Path, str],
|
| | py_logger: logging.Logger,
|
| | *,
|
| | is_rank0: bool = False,
|
| | enable_email: bool = False):
|
| | self.exp_id = exp_id
|
| | self.run_dir = Path(run_dir)
|
| | self.py_log = py_logger
|
| | self.email_sender = EmailSender(exp_id, enable=(is_rank0 and enable_email))
|
| | if is_rank0:
|
| | self.tb_log = SummaryWriter(run_dir)
|
| | else:
|
| | self.tb_log = None
|
| |
|
| |
|
| | try:
|
| | import git
|
| | repo = git.Repo(".")
|
| | git_info = str(repo.active_branch) + ' ' + str(repo.head.commit.hexsha)
|
| | except (ImportError, RuntimeError, TypeError):
|
| | print('Failed to fetch git info. Defaulting to None')
|
| | git_info = 'None'
|
| |
|
| | self.log_string('git', git_info)
|
| |
|
| |
|
| | job_id = os.environ.get('SLURM_JOB_ID', None)
|
| | if job_id is not None:
|
| | self.log_string('slurm_job_id', job_id)
|
| | self.email_sender.send(f'Job {job_id} started', f'Job started {run_dir}')
|
| |
|
| |
|
| | self.batch_timer: TimeEstimator = None
|
| | self.data_timer: PartialTimeEstimator = None
|
| |
|
| | self.nan_count = defaultdict(int)
|
| |
|
| | def log_scalar(self, tag: str, x: float, it: int):
|
| | if self.tb_log is None:
|
| | return
|
| | if math.isnan(x) and 'grad_norm' not in tag:
|
| | self.nan_count[tag] += 1
|
| | if self.nan_count[tag] == 10:
|
| | self.email_sender.send(
|
| | f'Nan detected in {tag} @ {self.run_dir}',
|
| | f'Nan detected in {tag} at iteration {it}; run_dir: {self.run_dir}')
|
| | else:
|
| | self.nan_count[tag] = 0
|
| | self.tb_log.add_scalar(tag, x, it)
|
| |
|
| | def log_metrics(self,
|
| | prefix: str,
|
| | metrics: dict[str, float],
|
| | it: int,
|
| | ignore_timer: bool = False):
|
| | msg = f'{self.exp_id}-{prefix} - it {it:6d}: '
|
| | metrics_msg = ''
|
| | for k, v in sorted(metrics.items()):
|
| | self.log_scalar(f'{prefix}/{k}', v, it)
|
| | metrics_msg += f'{k: >10}:{v:.7f},\t'
|
| |
|
| | if self.batch_timer is not None and not ignore_timer:
|
| | self.batch_timer.update()
|
| | avg_time = self.batch_timer.get_and_reset_avg_time()
|
| | data_time = self.data_timer.get_and_reset_avg_time()
|
| |
|
| |
|
| | self.log_scalar(f'{prefix}/avg_time', avg_time, it)
|
| | self.log_scalar(f'{prefix}/data_time', data_time, it)
|
| |
|
| | est = self.batch_timer.get_est_remaining(it)
|
| | est = datetime.timedelta(seconds=est)
|
| | if est.days > 0:
|
| | remaining_str = f'{est.days}d {est.seconds // 3600}h'
|
| | else:
|
| | remaining_str = f'{est.seconds // 3600}h {(est.seconds%3600) // 60}m'
|
| | eta = datetime.datetime.now(timezone(my_timezone)) + est
|
| | eta_str = eta.strftime('%Y-%m-%d %H:%M:%S %Z%z')
|
| | time_msg = f'avg_time:{avg_time:.3f},data:{data_time:.3f},remaining:{remaining_str},eta:{eta_str},\t'
|
| | msg = f'{msg} {time_msg}'
|
| |
|
| | msg = f'{msg} {metrics_msg}'
|
| | self.py_log.info(msg)
|
| |
|
| | def log_histogram(self, tag: str, hist: torch.Tensor, it: int):
|
| | if self.tb_log is None:
|
| | return
|
| |
|
| | hist = hist.cpu().numpy()
|
| | fig, ax = plt.subplots()
|
| | x_range = np.linspace(0, 1, len(hist))
|
| | ax.bar(x_range, hist, width=1 / (len(hist) - 1))
|
| | ax.set_xticks(x_range)
|
| | ax.set_xticklabels(x_range)
|
| | plt.tight_layout()
|
| | self.tb_log.add_figure(tag, fig, it)
|
| | plt.close()
|
| |
|
| | def log_image(self, prefix: str, tag: str, image: np.ndarray, it: int):
|
| | image_dir = self.run_dir / f'{prefix}_images'
|
| | image_dir.mkdir(exist_ok=True, parents=True)
|
| |
|
| | image = Image.fromarray(image)
|
| | image.save(image_dir / f'{it:09d}_{tag}.png')
|
| |
|
| | def log_audio(self,
|
| | prefix: str,
|
| | tag: str,
|
| | waveform: torch.Tensor,
|
| | it: Optional[int] = None,
|
| | *,
|
| | subdir: Optional[Path] = None,
|
| | sample_rate: int = 16000) -> Path:
|
| | if subdir is None:
|
| | audio_dir = self.run_dir / prefix
|
| | else:
|
| | audio_dir = self.run_dir / subdir / prefix
|
| | audio_dir.mkdir(exist_ok=True, parents=True)
|
| |
|
| | if it is None:
|
| | name = f'{tag}.flac'
|
| | else:
|
| | name = f'{it:09d}_{tag}.flac'
|
| |
|
| | torchaudio.save(audio_dir / name,
|
| | waveform.cpu().float(),
|
| | sample_rate=sample_rate,
|
| | channels_first=True)
|
| | return Path(audio_dir)
|
| |
|
| | def log_spectrogram(
|
| | self,
|
| | prefix: str,
|
| | tag: str,
|
| | spec: torch.Tensor,
|
| | it: Optional[int],
|
| | *,
|
| | subdir: Optional[Path] = None,
|
| | ):
|
| | if subdir is None:
|
| | spec_dir = self.run_dir / prefix
|
| | else:
|
| | spec_dir = self.run_dir / subdir / prefix
|
| | spec_dir.mkdir(exist_ok=True, parents=True)
|
| |
|
| | if it is None:
|
| | name = f'{tag}.png'
|
| | else:
|
| | name = f'{it:09d}_{tag}.png'
|
| |
|
| | plot_spectrogram(spec.cpu().float())
|
| | plt.tight_layout()
|
| | plt.savefig(spec_dir / name)
|
| | plt.close()
|
| |
|
| | def log_string(self, tag: str, x: str):
|
| | self.py_log.info(f'{tag} - {x}')
|
| | if self.tb_log is None:
|
| | return
|
| | self.tb_log.add_text(tag, x)
|
| |
|
| | def debug(self, x):
|
| | self.py_log.debug(x)
|
| |
|
| | def info(self, x):
|
| | self.py_log.info(x)
|
| |
|
| | def warning(self, x):
|
| | self.py_log.warning(x)
|
| |
|
| | def error(self, x):
|
| | self.py_log.error(x)
|
| |
|
| | def critical(self, x):
|
| | self.py_log.critical(x)
|
| |
|
| | self.email_sender.send(f'Error occurred in {self.run_dir}', x)
|
| |
|
| | def complete(self):
|
| | self.email_sender.send(f'Job completed in {self.run_dir}', 'Job completed')
|
| |
|
