72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
|
"""
|
||
|
|
A logger that queries CuPy metrics and passes that information to downstream loggers.
|
||
|
|
"""
|
||
|
|
from typing import Dict, Any, Optional, IO
|
||
|
|
import sys
|
||
|
|
|
||
|
|
from spacy import Language
|
||
|
|
from thinc.backends import context_pools
|
||
|
|
from thinc.util import has_cupy_gpu
|
||
|
|
from .util import LoggerT
|
||
|
|
|
||
|
|
|
||
|
|
def cupy_logger_v1(
|
||
|
|
prefix: str = "cupy",
|
||
|
|
) -> LoggerT:
|
||
|
|
"""Creates a logger that queries CuPy metrics and passes that information to downstream loggers.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
prefix (str):
|
||
|
|
All metric names are prefixed with this string using dot notation, e.g: `<prefix>.<metric>`.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
LoggerT: Logger instance.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
import cupy
|
||
|
|
except ImportError:
|
||
|
|
raise ImportError(
|
||
|
|
"The 'cupy' library could not be found - did you install it? "
|
||
|
|
"Alternatively, specify the 'ConsoleLogger' in the "
|
||
|
|
"'training.logger' config section, instead of the 'CuPyLogger'."
|
||
|
|
)
|
||
|
|
|
||
|
|
if not has_cupy_gpu:
|
||
|
|
raise ValueError(
|
||
|
|
"The 'cupy' library could be loaded but GPU support couldn't be initialized - "
|
||
|
|
"do you have the necessary CUDA libraries installed?"
|
||
|
|
)
|
||
|
|
|
||
|
|
def setup_logger(nlp: Language, stdout: IO = sys.stdout, stderr: IO = sys.stderr):
|
||
|
|
def to_mib(bytes: int) -> float:
|
||
|
|
return bytes / (1024.0**2)
|
||
|
|
|
||
|
|
def log_step(info: Optional[Dict[str, Any]]):
|
||
|
|
if info is None:
|
||
|
|
return
|
||
|
|
|
||
|
|
# We cannot rely on cupy.cuda.get_default_memory_pool() as thinc can
|
||
|
|
# use its own allocator/pool for interop with PyTorch/TensorFlow.
|
||
|
|
cupy_memory_pool = cupy.cuda.get_allocator().__self__
|
||
|
|
thinc_pool = context_pools.get()
|
||
|
|
|
||
|
|
if thinc_pool.get("pytorch") == cupy_memory_pool:
|
||
|
|
pool_source = "pytorch"
|
||
|
|
elif thinc_pool.get("tensorflow") == cupy_memory_pool:
|
||
|
|
pool_source = "tensorflow"
|
||
|
|
else:
|
||
|
|
pool_source = "default"
|
||
|
|
|
||
|
|
info[f"{prefix}.pool.source"] = pool_source
|
||
|
|
info[f"{prefix}.pool.acquired_mib"] = to_mib(cupy_memory_pool.total_bytes())
|
||
|
|
info[f"{prefix}.pool.used_mib"] = to_mib(cupy_memory_pool.used_bytes())
|
||
|
|
info[f"{prefix}.pool.free_mib"] = to_mib(cupy_memory_pool.free_bytes())
|
||
|
|
info[f"{prefix}.pool.num_free_blocks"] = cupy_memory_pool.n_free_blocks()
|
||
|
|
|
||
|
|
def finalize():
|
||
|
|
pass
|
||
|
|
|
||
|
|
return log_step, finalize
|
||
|
|
|
||
|
|
return setup_logger
|