| """ |
| Dumps things to tensorboard and console |
| """ |
|
|
| import os |
| import logging |
| import datetime |
| from typing import Dict |
| import numpy as np |
| from PIL import Image |
|
|
| from torch.utils.tensorboard import SummaryWriter |
| from tracker.utils.time_estimator import TimeEstimator |
|
|
|
|
| def tensor_to_numpy(image): |
| image_np = (image.numpy() * 255).astype('uint8') |
| return image_np |
|
|
|
|
| def detach_to_cpu(x): |
| return x.detach().cpu() |
|
|
|
|
| def fix_width_trunc(x): |
| return ('{:.9s}'.format('{:0.9f}'.format(x))) |
|
|
|
|
| class TensorboardLogger: |
| def __init__(self, run_dir, py_logger: logging.Logger, *, enabled_tb): |
| self.run_dir = run_dir |
| self.py_log = py_logger |
| if enabled_tb: |
| 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): |
| print('Failed to fetch git info. Defaulting to None') |
| git_info = 'None' |
|
|
| self.log_string('git', git_info) |
|
|
| |
| self.time_estimator: TimeEstimator = None |
|
|
| def log_scalar(self, tag, x, it): |
| if self.tb_log is None: |
| return |
| self.tb_log.add_scalar(tag, x, it) |
|
|
| def log_metrics(self, exp_id, prefix, metrics: Dict, it): |
| msg = f'{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.time_estimator is not None: |
| self.time_estimator.update() |
| avg_time = self.time_estimator.get_and_reset_avg_time() |
| est = self.time_estimator.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() + est |
| eta_str = eta.strftime('%Y-%m-%d %H:%M:%S') |
| time_msg = f'avg_time:{avg_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_image(self, stage_name, tag, image, it): |
| image_dir = os.path.join(self.run_dir, f'{stage_name}_images') |
| os.makedirs(image_dir, exist_ok=True) |
|
|
| image = Image.fromarray(image) |
| image.save(os.path.join(image_dir, f'{tag}_{it}.png')) |
|
|
| def log_string(self, tag, x): |
| 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) |
|
|
