feat(utils): 配置继承、注册表、日志与 Nvidia GPU 守卫
- config: yaml `_base_` 继承 + 深度合并 + 命令行 `--set a.b.c=value` 覆盖，检测循环引用 - registry: 组件按名注册与构建，模型/数据集共用 - device: require_cuda() 守卫，按 CLAUDE.md 约定禁止在无 Nvidia GPU 的设备上训练； 评估与导出走 resolve_device() 允许 CPU - labels: multilabel / multiclass 两种图像级标签形式的共用定义 - logging / seed: 控制台+文件日志、随机种子与 cudnn 确定性 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
0645fc19f7
commit
f4eff8beee
@@ -0,0 +1,23 @@
|
||||
"""通用工具：配置、注册表、日志、随机种子、设备检查。"""
|
||||
|
||||
from utils.config import Config, load_config
|
||||
from utils.device import require_cuda, resolve_device
|
||||
from utils.labels import LABEL_MODES, MULTICLASS, MULTILABEL, check_label_mode
|
||||
from utils.logging import get_logger, setup_logging
|
||||
from utils.registry import Registry
|
||||
from utils.seed import seed_everything
|
||||
|
||||
__all__ = [
|
||||
"Config",
|
||||
"load_config",
|
||||
"require_cuda",
|
||||
"resolve_device",
|
||||
"LABEL_MODES",
|
||||
"MULTILABEL",
|
||||
"MULTICLASS",
|
||||
"check_label_mode",
|
||||
"get_logger",
|
||||
"setup_logging",
|
||||
"Registry",
|
||||
"seed_everything",
|
||||
]
|
||||
+118
@@ -0,0 +1,118 @@
|
||||
"""YAML 配置加载：支持 `_base_` 继承、深度合并与命令行点号覆盖。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
from collections.abc import Iterable, Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
BASE_KEY = "_base_"
|
||||
|
||||
|
||||
class Config(dict):
|
||||
"""行为与 dict 一致，额外支持 `cfg.train.lr` 式属性访问。"""
|
||||
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
for key, value in list(self.items()):
|
||||
if isinstance(value, Mapping) and not isinstance(value, Config):
|
||||
self[key] = Config(value)
|
||||
|
||||
def __getattr__(self, item: str) -> Any:
|
||||
try:
|
||||
return self[item]
|
||||
except KeyError as exc: # pragma: no cover - 仅用于更友好的报错
|
||||
raise AttributeError(item) from exc
|
||||
|
||||
def get_path(self, dotted: str, default: Any = None) -> Any:
|
||||
node: Any = self
|
||||
for part in dotted.split("."):
|
||||
if not isinstance(node, Mapping) or part not in node:
|
||||
return default
|
||||
node = node[part]
|
||||
return node
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
def unwrap(value: Any) -> Any:
|
||||
if isinstance(value, Mapping):
|
||||
return {k: unwrap(v) for k, v in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [unwrap(v) for v in value]
|
||||
return value
|
||||
|
||||
return unwrap(self)
|
||||
|
||||
def dump(self, path: str | Path) -> None:
|
||||
Path(path).write_text(
|
||||
yaml.safe_dump(self.to_dict(), allow_unicode=True, sort_keys=False),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def deep_merge(base: Mapping, override: Mapping) -> dict:
|
||||
"""递归合并；override 中的标量/列表直接覆盖 base。"""
|
||||
merged = dict(base)
|
||||
for key, value in override.items():
|
||||
if key in merged and isinstance(merged[key], Mapping) and isinstance(value, Mapping):
|
||||
merged[key] = deep_merge(merged[key], value)
|
||||
else:
|
||||
merged[key] = value
|
||||
return merged
|
||||
|
||||
|
||||
def _load_raw(path: Path, seen: set[Path]) -> dict:
|
||||
path = path.resolve()
|
||||
if path in seen:
|
||||
raise ValueError(f"配置 `{BASE_KEY}` 出现循环引用: {path}")
|
||||
seen.add(path)
|
||||
|
||||
data = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
|
||||
if not isinstance(data, dict):
|
||||
raise TypeError(f"配置文件必须是映射: {path}")
|
||||
|
||||
bases = data.pop(BASE_KEY, [])
|
||||
if isinstance(bases, str):
|
||||
bases = [bases]
|
||||
|
||||
merged: dict = {}
|
||||
for base in bases:
|
||||
base_path = (path.parent / base).resolve()
|
||||
merged = deep_merge(merged, _load_raw(base_path, set(seen)))
|
||||
return deep_merge(merged, data)
|
||||
|
||||
|
||||
def parse_override(item: str) -> tuple[str, Any]:
|
||||
"""解析 `train.lr=1e-3`，值按 Python 字面量尝试转换，失败则保留字符串。"""
|
||||
if "=" not in item:
|
||||
raise ValueError(f"覆盖项格式应为 key=value，收到: {item}")
|
||||
key, raw = item.split("=", 1)
|
||||
try:
|
||||
value = ast.literal_eval(raw)
|
||||
except (ValueError, SyntaxError):
|
||||
value = raw
|
||||
return key.strip(), value
|
||||
|
||||
|
||||
def apply_overrides(cfg: dict, overrides: Iterable[str]) -> dict:
|
||||
for item in overrides:
|
||||
key, value = parse_override(item)
|
||||
parts = key.split(".")
|
||||
node = cfg
|
||||
for part in parts[:-1]:
|
||||
nxt = node.get(part)
|
||||
if not isinstance(nxt, dict):
|
||||
nxt = {}
|
||||
node[part] = nxt
|
||||
node = nxt
|
||||
node[parts[-1]] = value
|
||||
return cfg
|
||||
|
||||
|
||||
def load_config(path: str | Path, overrides: Iterable[str] = ()) -> Config:
|
||||
cfg = _load_raw(Path(path), set())
|
||||
cfg = apply_overrides(cfg, overrides)
|
||||
cfg.setdefault("config_path", str(Path(path).resolve()))
|
||||
return Config(cfg)
|
||||
@@ -0,0 +1,54 @@
|
||||
"""设备检查。
|
||||
|
||||
项目硬性约定（见 CLAUDE.MD）：禁止在没有 Nvidia GPU 的设备上训练。
|
||||
所有训练入口必须先调用 `require_cuda()`，检查不通过直接抛错退出。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class NoNvidiaGpuError(RuntimeError):
|
||||
"""当前设备不具备可用的 Nvidia GPU。"""
|
||||
|
||||
|
||||
def cuda_devices() -> list[str]:
|
||||
"""返回可用的 Nvidia GPU 名称列表；没有 torch 或没有 GPU 时返回空列表。"""
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
return []
|
||||
if not torch.cuda.is_available() or torch.version.cuda is None:
|
||||
return []
|
||||
return [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]
|
||||
|
||||
|
||||
def require_cuda() -> list[str]:
|
||||
"""训练前的守卫；无 Nvidia GPU 时抛出 `NoNvidiaGpuError`。"""
|
||||
try:
|
||||
import torch
|
||||
except ImportError as exc:
|
||||
raise NoNvidiaGpuError("未安装 PyTorch，无法确认 Nvidia GPU，训练已中止。") from exc
|
||||
|
||||
if torch.version.cuda is None:
|
||||
raise NoNvidiaGpuError(
|
||||
f"当前 PyTorch ({torch.__version__}) 不是 CUDA 构建（可能是 CPU/ROCm/MPS 版本），"
|
||||
"本项目禁止在非 Nvidia GPU 设备上训练。"
|
||||
)
|
||||
if not torch.cuda.is_available() or torch.cuda.device_count() == 0:
|
||||
raise NoNvidiaGpuError(
|
||||
"未检测到可用的 Nvidia GPU，本项目禁止在无 GPU 的设备上训练。"
|
||||
"评估/推理请使用 eval.py 并显式指定 device=cpu。"
|
||||
)
|
||||
return cuda_devices()
|
||||
|
||||
|
||||
def resolve_device(device: str = "cuda"):
|
||||
"""把配置中的 device 字符串解析为 `torch.device`。
|
||||
|
||||
仅用于评估/推理等允许 CPU 的场景；训练路径请先走 `require_cuda()`。
|
||||
"""
|
||||
import torch
|
||||
|
||||
if device.startswith("cuda") and not torch.cuda.is_available():
|
||||
raise NoNvidiaGpuError(f"请求了 {device}，但当前设备没有可用的 Nvidia GPU。")
|
||||
return torch.device(device)
|
||||
@@ -0,0 +1,13 @@
|
||||
"""图像级标签形式的共用定义（模型与数据集都要用到，放这里避免互相 import）。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
MULTILABEL = "multilabel"
|
||||
MULTICLASS = "multiclass"
|
||||
LABEL_MODES = (MULTILABEL, MULTICLASS)
|
||||
|
||||
|
||||
def check_label_mode(label_mode: str) -> str:
|
||||
if label_mode not in LABEL_MODES:
|
||||
raise ValueError(f"未知 label_mode: {label_mode}，可选 {LABEL_MODES}")
|
||||
return label_mode
|
||||
@@ -0,0 +1,36 @@
|
||||
"""日志：控制台 + 可选文件输出。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_FORMAT = "[%(asctime)s][%(levelname)s][%(name)s] %(message)s"
|
||||
_DATEFMT = "%Y-%m-%d %H:%M:%S"
|
||||
|
||||
|
||||
def setup_logging(log_file: str | Path | None = None, level: int | str = logging.INFO) -> logging.Logger:
|
||||
root = logging.getLogger("mte")
|
||||
root.setLevel(level)
|
||||
root.handlers.clear()
|
||||
root.propagate = False
|
||||
|
||||
formatter = logging.Formatter(_FORMAT, datefmt=_DATEFMT)
|
||||
|
||||
stream = logging.StreamHandler(sys.stdout)
|
||||
stream.setFormatter(formatter)
|
||||
root.addHandler(stream)
|
||||
|
||||
if log_file is not None:
|
||||
path = Path(log_file)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
file_handler = logging.FileHandler(path, encoding="utf-8")
|
||||
file_handler.setFormatter(formatter)
|
||||
root.addHandler(file_handler)
|
||||
|
||||
return root
|
||||
|
||||
|
||||
def get_logger(name: str = "mte") -> logging.Logger:
|
||||
return logging.getLogger(name if name.startswith("mte") else f"mte.{name}")
|
||||
@@ -0,0 +1,50 @@
|
||||
"""极简注册表：按名字注册 / 构建模型、数据集等组件。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Iterator
|
||||
from typing import Any, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class Registry:
|
||||
def __init__(self, name: str) -> None:
|
||||
self.name = name
|
||||
self._entries: dict[str, Callable[..., Any]] = {}
|
||||
|
||||
def register(self, name: str | None = None) -> Callable[[T], T]:
|
||||
"""装饰器；`name` 缺省时使用被装饰对象的 __name__。"""
|
||||
|
||||
def wrapper(obj: T) -> T:
|
||||
key = name or getattr(obj, "__name__", None)
|
||||
if key is None:
|
||||
raise ValueError(f"无法推断注册名: {obj!r}")
|
||||
if key in self._entries:
|
||||
raise KeyError(f"{self.name} 中已存在同名条目: {key}")
|
||||
self._entries[key] = obj # type: ignore[assignment]
|
||||
return obj
|
||||
|
||||
return wrapper
|
||||
|
||||
def get(self, name: str) -> Callable[..., Any]:
|
||||
if name not in self._entries:
|
||||
raise KeyError(f"{self.name} 中找不到 '{name}'，可用: {sorted(self._entries)}")
|
||||
return self._entries[name]
|
||||
|
||||
def build(self, cfg: dict, **kwargs: Any) -> Any:
|
||||
"""按 `{"name": ..., 其余作为构造参数}` 构建对象。"""
|
||||
cfg = dict(cfg)
|
||||
name = cfg.pop("name", None)
|
||||
if name is None:
|
||||
raise KeyError(f"构建 {self.name} 需要 'name' 字段: {cfg}")
|
||||
return self.get(name)(**{**cfg, **kwargs})
|
||||
|
||||
def __contains__(self, name: object) -> bool:
|
||||
return name in self._entries
|
||||
|
||||
def __iter__(self) -> Iterator[str]:
|
||||
return iter(self._entries)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._entries)
|
||||
@@ -0,0 +1,33 @@
|
||||
"""随机种子与 cudnn 确定性设置。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import random
|
||||
|
||||
|
||||
def seed_everything(seed: int, deterministic: bool = False) -> int:
|
||||
random.seed(seed)
|
||||
os.environ["PYTHONHASHSEED"] = str(seed)
|
||||
|
||||
try:
|
||||
import numpy as np
|
||||
|
||||
np.random.seed(seed)
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
if deterministic:
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
else:
|
||||
torch.backends.cudnn.benchmark = True
|
||||
except ImportError: # pragma: no cover
|
||||
pass
|
||||
|
||||
return seed
|
||||
Reference in New Issue
Block a user