|
|
|
|
|
|
|
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) |
|
|
|
|
|
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 |
|
|
|
|
|
|
|
tb_logger = MyTrainingLogger() |
|
|
|
|
|
|
|
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 |
|
|
