Spaces:
Running
on
A10G
Running
on
A10G
File size: 3,144 Bytes
320e465 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
"""
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
# Get current git info for logging
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)
# used when logging metrics
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)
|