File size: 3,841 Bytes
7f316fe | 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 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the CC-by-NC license found in the
# LICENSE file in the root directory of this source tree.
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.")
|