| |
| |
| |
| |
| |
|
|
| import logging |
| import os |
| from logging import Logger |
| from pathlib import Path |
| from typing import Optional |
|
|
| import torch |
| import wandb |
| from omegaconf import OmegaConf |
|
|
|
|
| def get_logger(log_path: str, rank: int): |
| if rank != 0: |
| return logging.getLogger("dummy") |
|
|
| logger = logging.getLogger() |
| default_level = logging.INFO |
|
|
| if logger.hasHandlers(): |
| logger.handlers.clear() |
|
|
| logger.setLevel(default_level) |
|
|
| formatter = logging.Formatter( |
| "%(levelname)s | %(asctime)s | %(message)s", "%Y-%m-%d %H:%M:%S" |
| ) |
|
|
| info_file_handler = logging.FileHandler(log_path, mode="a") |
| info_file_handler.setLevel(default_level) |
| info_file_handler.setFormatter(formatter) |
| logger.addHandler(info_file_handler) |
|
|
| console_handler = logging.StreamHandler() |
| console_handler.setLevel(default_level) |
| console_handler.setFormatter(formatter) |
| logger.addHandler(console_handler) |
|
|
| return logger |
|
|
|
|
| class TrainLogger: |
| def __init__(self, log_dir: Path, rank: int, cfg: bool = False): |
| self.log_dir = log_dir |
| self.cfg = cfg |
|
|
| self._init_text_logger(rank=rank) |
|
|
| self.enable_wandb = self.cfg.logging.enable_wandb and (rank == 0) |
|
|
| if self.enable_wandb: |
| self._init_wandb() |
|
|
| def _init_text_logger(self, rank: int): |
| log_path = self.log_dir / self.cfg.logging.log_file_name |
| self._logger = get_logger(log_path=log_path, rank=rank) |
|
|
| def _init_wandb( |
| self, |
| ): |
| wandb_run_id_path = self.log_dir / "wandb_run.id" |
|
|
| try: |
| wandb_run_id = wandb_run_id_path.read_text() |
| except FileNotFoundError: |
| wandb_run_id = wandb.util.generate_id() |
| wandb_run_id_path.write_text(wandb_run_id) |
|
|
| self.wandb_logger = wandb.init( |
| id=wandb_run_id, |
| project=self.cfg.logging.project, |
| group=self.cfg.logging.group, |
| dir=self.log_dir, |
| entity=self.cfg.logging.entity, |
| resume="allow", |
| config=OmegaConf.to_container(self.cfg, resolve=True), |
| ) |
|
|
| def log_metric(self, value: float, name: str, stage: bool, step: int) -> None: |
| self._logger.info(f"[{step}] {stage} {name}: {value:.3f}") |
|
|
| if self.enable_wandb: |
| self.wandb_logger.log(data={f"{stage}/{name}": value}, step=step) |
|
|
| def log_lr(self, value: float, step: int) -> None: |
| if self.enable_wandb: |
| self.wandb_logger.log(data={"Optimization/LR": value}, step=step) |
|
|
| def info(self, msg: str, step: Optional[int] = None) -> None: |
| step_str = f"[{step}] " if step else "" |
| self._logger.info(f"{step_str}{msg}") |
|
|
| def warning(self, msg: str) -> None: |
| self._logger.warning(msg) |
|
|
| def finish(self) -> None: |
| for handler in self._logger.handlers: |
| if isinstance(handler, logging.FileHandler): |
| handler.close() |
|
|
| if self.enable_wandb: |
| wandb.finish() |
|
|
| @staticmethod |
| def log_devices(device: torch.device, logger: Logger) -> None: |
| if device.type == "cuda": |
| logger.info("Found {} CUDA devices.".format(torch.cuda.device_count())) |
| for i in range(torch.cuda.device_count()): |
| props = torch.cuda.get_device_properties(i) |
| logger.info( |
| "{} \t Memory: {:.2f}GB".format( |
| props.name, props.total_memory / (1024**3) |
| ) |
| ) |
| else: |
| logger.warning("WARNING: Using device {}".format(device)) |
| logger.info(f"Found {os.cpu_count()} total number of CPUs.") |
|
|
