diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..2fc0559 --- /dev/null +++ b/utils/__init__.py @@ -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", +] diff --git a/utils/config.py b/utils/config.py new file mode 100644 index 0000000..1e619b3 --- /dev/null +++ b/utils/config.py @@ -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) diff --git a/utils/device.py b/utils/device.py new file mode 100644 index 0000000..fa3f59c --- /dev/null +++ b/utils/device.py @@ -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) diff --git a/utils/labels.py b/utils/labels.py new file mode 100644 index 0000000..753bd33 --- /dev/null +++ b/utils/labels.py @@ -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 diff --git a/utils/logging.py b/utils/logging.py new file mode 100644 index 0000000..b5fb539 --- /dev/null +++ b/utils/logging.py @@ -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}") diff --git a/utils/registry.py b/utils/registry.py new file mode 100644 index 0000000..62a25ed --- /dev/null +++ b/utils/registry.py @@ -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) diff --git a/utils/seed.py b/utils/seed.py new file mode 100644 index 0000000..745b549 --- /dev/null +++ b/utils/seed.py @@ -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