File size: 2,874 Bytes
864ec44 |
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 |
# Author: Bingxin Ke
# Last modified: 2024-03-12
import logging
import os
import sys
import wandb
from tabulate import tabulate
from torch.utils.tensorboard import SummaryWriter
def config_logging(cfg_logging, out_dir=None):
file_level = cfg_logging.get("file_level", 10)
console_level = cfg_logging.get("console_level", 10)
log_formatter = logging.Formatter(cfg_logging["format"])
root_logger = logging.getLogger()
root_logger.handlers.clear()
root_logger.setLevel(min(file_level, console_level))
if out_dir is not None:
_logging_file = os.path.join(
out_dir, cfg_logging.get("filename", "logging.log")
)
file_handler = logging.FileHandler(_logging_file)
file_handler.setFormatter(log_formatter)
file_handler.setLevel(file_level)
root_logger.addHandler(file_handler)
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(log_formatter)
console_handler.setLevel(console_level)
root_logger.addHandler(console_handler)
# Avoid pollution by packages
logging.getLogger("PIL").setLevel(logging.INFO)
logging.getLogger("matplotlib").setLevel(logging.INFO)
class MyTrainingLogger:
"""Tensorboard + wandb logger"""
writer: SummaryWriter
is_initialized = False
def __init__(self) -> None:
pass
def set_dir(self, tb_log_dir):
if self.is_initialized:
raise ValueError("Do not initialize writer twice")
self.writer = SummaryWriter(tb_log_dir)
self.is_initialized = True
def log_dic(self, scalar_dic, global_step, walltime=None):
for k, v in scalar_dic.items():
self.writer.add_scalar(k, v, global_step=global_step, walltime=walltime)
return
# global instance
tb_logger = MyTrainingLogger()
# -------------- wandb tools --------------
def init_wandb(enable: bool, **kwargs):
if enable:
run = wandb.init(sync_tensorboard=True, **kwargs)
else:
run = wandb.init(mode="disabled")
return run
def log_slurm_job_id(step):
global tb_logger
_jobid = os.getenv("SLURM_JOB_ID")
if _jobid is None:
_jobid = -1
tb_logger.writer.add_scalar("job_id", int(_jobid), global_step=step)
logging.debug(f"Slurm job_id: {_jobid}")
def load_wandb_job_id(out_dir):
with open(os.path.join(out_dir, "WANDB_ID"), "r") as f:
wandb_id = f.read()
return wandb_id
def save_wandb_job_id(run, out_dir):
with open(os.path.join(out_dir, "WANDB_ID"), "w+") as f:
f.write(run.id)
def eval_dic_to_text(val_metrics: dict, dataset_name: str, sample_list_path: str):
eval_text = f"Evaluation metrics:\n\
on dataset: {dataset_name}\n\
over samples in: {sample_list_path}\n"
eval_text += tabulate([val_metrics.keys(), val_metrics.values()])
return eval_text
|