"""
A logger that queries PyTorch metrics and passes that information to downstream loggers.
"""
from typing import Dict, Any, Optional, Tuple, IO
import re
import sys
from spacy import Language
from .util import LoggerT
def pytorch_logger_v1(
prefix: str = "pytorch",
device: int = 0,
cuda_mem_pool: str = "all",
cuda_mem_metric: str = "all",
) -> LoggerT:
try:
import torch
except ImportError:
raise ImportError(
"The 'torch' library could not be found - did you install it? "
"Alternatively, specify the 'ConsoleLogger' in the "
"'training.logger' config section, instead of the 'PyTorchLogger'."
)
def setup_logger(nlp: Language, stdout: IO = sys.stdout, stderr: IO = sys.stderr):
expected_cuda_mem_pool = ("all", "large_pool", "small_pool")
expected_cuda_mem_metric = ("all", "current", "peak", "allocated", "free")
if cuda_mem_pool not in expected_cuda_mem_pool:
raise ValueError(
f"Got CUDA memory pool '{cuda_mem_pool}', but expected one of: '{expected_cuda_mem_pool}'"
)
elif cuda_mem_metric not in expected_cuda_mem_metric:
raise ValueError(
f"Got CUDA memory metric '{cuda_mem_metric}', but expected one of: '{expected_cuda_mem_metric}'"
)
def normalize_mem_value_to_mb(name: str, value: int) -> Tuple[str, float]:
if "_bytes" in name:
return re.sub("_bytes", "_megabytes", name), value / (1024.0**2)
else:
return name, value
def log_step(info: Optional[Dict[str, Any]]):
if info is None:
return
cuda_mem_stats = torch.cuda.memory_stats(device)
for stat, val in cuda_mem_stats.items():
splits = stat.split(".")
if len(splits) == 3:
name, pool, metric = splits
name, val = normalize_mem_value_to_mb(name, val)
if pool != cuda_mem_pool:
continue
elif cuda_mem_metric != "all" and metric != cuda_mem_metric:
continue
info[f"{prefix}.{name}.{pool}.{metric}"] = val
elif len(splits) == 2:
name, metric = splits
name, val = normalize_mem_value_to_mb(name, val)
if cuda_mem_metric != "all" and metric != cuda_mem_metric:
continue
info[f"{prefix}.{name}.{metric}"] = val
else:
# Either global statistic or something that we haven't accounted for,
# e.g: a newly added statistic. So, we'll just include it to be safe.
info[f"{prefix}.{stat}"] = val
def finalize():
pass
return log_step, finalize
return setup_logger