Spaces:
Sleeping
Sleeping
File size: 2,294 Bytes
2875fe6 |
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 |
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import time
import torch
from utils.io_utils import write_args, save_config_to_yaml
class Logger(object):
def __init__(self, args):
self.args = args
self.save_dir = args.save_dir
os.makedirs(self.save_dir, exist_ok=True)
# save the args and config
self.config_dir = os.path.join(self.save_dir, "configs")
os.makedirs(self.config_dir, exist_ok=True)
file_name = os.path.join(self.config_dir, "args.txt")
write_args(args, file_name)
log_dir = os.path.join(self.save_dir, "logs")
if not os.path.exists(log_dir):
os.makedirs(log_dir, exist_ok=True)
self.text_writer = open(os.path.join(log_dir, "log.txt"), "a") # 'w')
if args.tensorboard:
self.log_info("using tensorboard")
self.tb_writer = torch.utils.tensorboard.SummaryWriter(
log_dir=log_dir
) # tensorboard.SummaryWriter(log_dir=log_dir)
else:
self.tb_writer = None
def save_config(self, config):
save_config_to_yaml(config, os.path.join(self.config_dir, "config.yaml"))
def log_info(self, info, check_primary=True):
print(info)
info = str(info)
time_str = time.strftime("%Y-%m-%d-%H-%M")
info = "{}: {}".format(time_str, info)
if not info.endswith("\n"):
info += "\n"
self.text_writer.write(info)
self.text_writer.flush()
def add_scalar(self, **kargs):
"""Log a scalar variable."""
if self.tb_writer is not None:
self.tb_writer.add_scalar(**kargs)
def add_scalars(self, **kargs):
"""Log a scalar variable."""
if self.tb_writer is not None:
self.tb_writer.add_scalars(**kargs)
def add_image(self, **kargs):
"""Log a scalar variable."""
if self.tb_writer is not None:
self.tb_writer.add_image(**kargs)
def add_images(self, **kargs):
"""Log a scalar variable."""
if self.tb_writer is not None:
self.tb_writer.add_images(**kargs)
def close(self):
self.text_writer.close()
self.tb_writer.close()
|