import os
import sys
import warnings
from collections import defaultdict
from copy import deepcopy
from typing import Dict, Optional, Tuple
import numpy as np
import torch
from loguru import logger
from torch.types import Number
from df_local.modules import GroupedLinearEinsum
from df_local.utils import get_branch_name, get_commit_hash, get_device, get_host
_logger_initialized = False
WARN_ONCE_NO = logger.level("WARNING").no + 1
DEPRECATED_NO = logger.level("WARNING").no + 2
def init_logger(file: Optional[str] = None, level: str = "INFO", model: Optional[str] = None):
global _logger_initialized, _duplicate_filter
if _logger_initialized:
logger.debug("Logger already initialized.")
else:
logger.remove()
level = level.upper()
if level.lower() != "none":
log_format = Formatter(debug=logger.level(level).no <= logger.level("DEBUG").no).format
logger.add(
sys.stdout,
level=level,
format=log_format,
filter=lambda r: r["level"].no not in {WARN_ONCE_NO, DEPRECATED_NO},
)
if file is not None:
logger.add(
file,
level=level,
format=log_format,
filter=lambda r: r["level"].no != WARN_ONCE_NO,
)
logger.info(f"Running on torch {torch.__version__}")
logger.info(f"Running on host {get_host()}")
commit = get_commit_hash()
if commit is not None:
logger.info(f"Git commit: {commit}, branch: {get_branch_name()}")
if (jobid := os.getenv("SLURM_JOB_ID")) is not None:
logger.info(f"Slurm jobid: {jobid}")
logger.level("WARNONCE", no=WARN_ONCE_NO, color="<yellow><bold>")
logger.add(
sys.stderr,
level=max(logger.level(level).no, WARN_ONCE_NO),
format=log_format,
filter=lambda r: r["level"].no == WARN_ONCE_NO and _duplicate_filter(r),
)
logger.level("DEPRECATED", no=DEPRECATED_NO, color="<yellow><bold>")
logger.add(
sys.stderr,
level=max(logger.level(level).no, DEPRECATED_NO),
format=log_format,
filter=lambda r: r["level"].no == DEPRECATED_NO and _duplicate_filter(r),
)
if model is not None:
logger.info("Loading model settings of {}", os.path.basename(model.rstrip("/")))
_logger_initialized = True
def warn_once(message, *args, **kwargs):
logger.log("WARNONCE", message, *args, **kwargs)
def log_deprecated(message, *args, **kwargs):
logger.log("DEPRECATED", message, *args, **kwargs)
class Formatter:
def __init__(self, debug=False):
if debug:
self.fmt = (
"<green>{time:YYYY-MM-DD HH:mm:ss}</green>"
" | <level>{level: <8}</level>"
" | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan>"
" | <level>{message}</level>"
)
else:
self.fmt = (
"<green>{time:YYYY-MM-DD HH:mm:ss}</green>"
" | <level>{level: <8}</level>"
" | <cyan>DF</cyan>"
" | <level>{message}</level>"
)
self.fmt += "\n{exception}"
def format(self, record):
if record["level"].no == WARN_ONCE_NO:
return self.fmt.replace("{level: <8}", "WARNING ")
return self.fmt
def _metrics_key(k_: Tuple[str, float]):
k0 = k_[0]
ks = k0.split("_")
if len(ks) > 2:
try:
return int(ks[-1])
except ValueError:
return 1000
elif k0 == "loss":
return -999
elif "loss" in k0.lower():
return -998
elif k0 == "lr":
return 998
elif k0 == "wd":
return 999
else:
return -101
def log_metrics(prefix: str, metrics: Dict[str, Number], level="INFO"):
msg = ""
stages = defaultdict(str)
loss_msg = ""
for n, v in sorted(metrics.items(), key=_metrics_key):
if abs(v) > 1e-3:
m = f" | {n}: {v:.5f}"
else:
m = f" | {n}: {v:.3E}"
if "stage" in n:
s = n.split("stage_")[1].split("_snr")[0]
stages[s] += m.replace(f"stage_{s}_", "")
elif ("valid" in prefix or "test" in prefix) and "loss" in n.lower():
loss_msg += m
else:
msg += m
for s, msg_s in stages.items():
logger.log(level, f"{prefix} | stage {s}" + msg_s)
if len(stages) == 0:
logger.log(level, prefix + msg)
if len(loss_msg) > 0:
logger.log(level, prefix + loss_msg)
class DuplicateFilter:
"""
Filters away duplicate log messages.
Modified version of: https://stackoverflow.com/a/60462619
"""
def __init__(self):
self.msgs = set()
def __call__(self, record) -> bool:
k = f"{record['level']}{record['message']}"
if k in self.msgs:
return False
else:
self.msgs.add(k)
return True
_duplicate_filter = DuplicateFilter()
def log_model_summary(model: torch.nn.Module, verbose=False):
try:
import ptflops
except ImportError:
logger.debug("Failed to import ptflops. Cannot print model summary.")
return
from df_local.model import ModelParams
# Generate input of 1 second audio
# Necessary inputs are:
# spec: [B, 1, T, F, 2], F: freq bin
# feat_erb: [B, 1, T, E], E: ERB bands
# feat_spec: [B, 2, T, C*2], C: Complex features
p = ModelParams()
b = 1
t = p.sr // p.hop_size
device = get_device()
spec = torch.randn([b, 1, t, p.fft_size // 2 + 1, 2]).to(device)
feat_erb = torch.randn([b, 1, t, p.nb_erb]).to(device)
feat_spec = torch.randn([b, 1, t, p.nb_df, 2]).to(device)
warnings.filterwarnings("ignore", "RNN module weights", category=UserWarning, module="torch")
macs, params = ptflops.get_model_complexity_info(
deepcopy(model),
(t,),
input_constructor=lambda _: {"spec": spec, "feat_erb": feat_erb, "feat_spec": feat_spec},
as_strings=False,
print_per_layer_stat=verbose,
verbose=verbose,
custom_modules_hooks={
GroupedLinearEinsum: grouped_linear_flops_counter_hook,
},
)
logger.info(f"Model complexity: {params/1e6:.3f}M #Params, {macs/1e6:.1f}M MACS")
def grouped_linear_flops_counter_hook(module: GroupedLinearEinsum, input, output):
# input: ([B, T, I],)
# output: [B, T, H]
input = input[0] # [B, T, I]
output_last_dim = module.weight.shape[-1]
input = input.unflatten(-1, (module.groups, module.ws)) # [B, T, G, I/G]
# GroupedLinear calculates "...gi,...gih->...gh"
weight_flops = np.prod(input.shape) * output_last_dim
module.__flops__ += int(weight_flops) # type: ignore