| 21 | |
| 22 | |
| 23 | class TensorboardLogger: |
| 24 | def __init__(self, place="./runs/", file_dir='./logs/', file_logger_name=None, random_id=True, tb_comment=None, |
| 25 | **tb_kwargs) -> None: |
| 26 | """tensorboard logger |
| 27 | |
| 28 | Args: |
| 29 | place (str, optional): place to save tb logging. Defaults to './tb_runs/'. |
| 30 | file_logger_name (str, optional): name of file logger. Defaults to None. |
| 31 | random_id: (bool, optional): append id after tensorboard dir. Defaults to True |
| 32 | tb_kwargs: kwargs for tensorboardX.writer.SummaryWriter |
| 33 | comment: Optional[str] = "", |
| 34 | purge_step: Optional[int] = None, |
| 35 | max_queue: Optional[int] = 10, |
| 36 | flush_secs: Optional[int] = 120, |
| 37 | filename_suffix: Optional[str] = '', |
| 38 | write_to_disk: Optional[bool] = True, |
| 39 | log_dir: Optional[str] = None, |
| 40 | comet_config: Optional[dict] = {"disabled": True}, |
| 41 | """ |
| 42 | if not place_exists(place): |
| 43 | os.mkdir(place) |
| 44 | print("Created directory: ", place) |
| 45 | |
| 46 | if random_id: |
| 47 | # use random id as dir name |
| 48 | id = generate_id(8) + '' if tb_comment is None else '_{}'.format(tb_comment) |
| 49 | place = os.path.join(place, id) |
| 50 | os.mkdir(place) |
| 51 | else: |
| 52 | # use time as dir name |
| 53 | stf_time = time.strftime("%m-%d %H:%M", time.localtime()) |
| 54 | name = stf_time + '' if tb_comment is None else '_{}'.format(tb_comment) |
| 55 | place = os.path.join(place, name) |
| 56 | os.mkdir(place) |
| 57 | |
| 58 | self.writer = writer.SummaryWriter(place, **tb_kwargs) |
| 59 | |
| 60 | assert file_logger_name is not None, '@file_logger_name should be set' |
| 61 | self.file_writer = PrintLogger(os.path.join(file_dir, file_logger_name + '.log')) |
| 62 | |
| 63 | self._tb_print = partial(self.file_writer.log, level=logging.DEBUG) |
| 64 | |
| 65 | def print(self, msg: str, level=logging.INFO): |
| 66 | self.file_writer.log(msg, level) |
| 67 | |
| 68 | @beartype |
| 69 | def log_scalar(self, tag: str, value: float, step: int): |
| 70 | self._tb_print(f"add tb scalar {tag}: {value}") |
| 71 | self.writer.add_scalar(tag, value, step) |
| 72 | |
| 73 | @beartype |
| 74 | def log_scalars(self, tag: str, values: dict, step: int, on_one_fig: bool = False): |
| 75 | self._tb_print(f"add tb scalars {tag}: {values}") |
| 76 | if not on_one_fig: |
| 77 | for t, v in values.items(): |
| 78 | self.writer.add_scalar(tag + '/' + t, v, step) |
| 79 | else: |
| 80 | self.writer.add_scalars(tag, values) |