diff --git a/README.md b/README.md index d7bf58a..b68ba77 100644 --- a/README.md +++ b/README.md @@ -89,6 +89,24 @@ uv run python scripts/export_onnx.py --config configs/train/baseline.yaml \ --checkpoint runs/baseline/best.pth --output runs/baseline/model.onnx ``` +换骨架用对应的配置对（训练与评估必须成对，否则构建出的骨架与权重对不上）： + +| 骨架 | 训练 | 评估 | work_dir | +| --- | --- | --- | --- | +| `tiny_conv`（默认） | `configs/train/baseline.yaml` | `configs/eval/baseline.yaml` | `runs/baseline` | +| `resnet18` | `configs/train/resnet18.yaml` | `configs/eval/resnet18.yaml` | `runs/resnet18` | + +```bash +uv run python train.py --config configs/train/resnet18.yaml +uv run python eval.py --config configs/eval/resnet18.yaml +uv run python scripts/export_onnx.py --config configs/train/resnet18.yaml \ + --checkpoint runs/resnet18/best.pth --output runs/resnet18/model.onnx +``` + +要加别的骨架（mobilenet / shufflenet），照 [configs/model/resnet18.yaml](configs/model/resnet18.yaml) +复制一份改 `arch` 即可——注意保留里面的 `_replace_: true`，否则 `tiny_conv` 的 +`widths/depths/block` 会残留下来传给新骨架。 + 产物写在 `work_dir`（默认 `runs//`）：`train.log`、`config.yaml`、`last.pth`、`best.pth`、`swanlog/`。 ## 多类 / 单类分类 diff --git a/configs/eval/resnet18.yaml b/configs/eval/resnet18.yaml new file mode 100644 index 0000000..3601819 --- /dev/null +++ b/configs/eval/resnet18.yaml @@ -0,0 +1,13 @@ +# resnet18 版本的评估配置。 +# 必须和训练时用的模型配置一致，否则构建出来的骨架与权重对不上。 +_base_: + - ../train/resnet18.yaml + +device: auto +checkpoint: runs/resnet18/best.pth + +eval: + split: val + score_threshold: 0.3 + iou_threshold: 0.5 + topk: 50 diff --git a/configs/model/resnet18.yaml b/configs/model/resnet18.yaml index ce67959..f2df69d 100644 --- a/configs/model/resnet18.yaml +++ b/configs/model/resnet18.yaml @@ -2,6 +2,9 @@ model: name: defect_net backbone: + # 整块替换：不加这行的话，从 tiny_conv 继承来的 widths/depths/block 会残留， + # 传给 TorchvisionBackbone 直接报错。见 utils/config.py::deep_merge + _replace_: true name: torchvision arch: resnet18 pretrained: true diff --git a/tests/test_config.py b/tests/test_config.py index 24525f4..ecb3666 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,3 +1,5 @@ +from pathlib import Path + import pytest pytest.importorskip("yaml") @@ -83,3 +85,52 @@ def test_registry_reports_incompatible_fields_actionably(): with pytest.raises(TypeError, match=r"不接受的字段.*widths"): BACKBONES.build({"name": "torchvision", "arch": "resnet18", "widths": [32]}) + + +@pytest.mark.parametrize( + "config_path", + sorted(str(p) for p in Path("configs").rglob("*.yaml")), +) +def test_every_shipped_config_builds_its_components(config_path): + """仓库里每个配置都必须能真的构建出组件。 + + 回归测试：configs/train/resnet18.yaml 曾经无法使用——它同时继承 tiny_conv 与 + resnet18 的模型配置，合并后 backbone 混了两边的字段，构建时直接报错。 + """ + torch = pytest.importorskip("torch") + from datasets.yolo import load_classes + from models import build_model + from models.registry import BACKBONES + + cfg = load_config(config_path) + + if "model" in cfg: + if "classes" not in cfg: + # configs/model/*.yaml 是片段，只能被 _base_ 组合，单独用没有 classes + pytest.skip(f"{config_path} 是配置片段，不含 classes") + model_cfg = dict(cfg["model"]) + model_cfg.setdefault("num_classes", len(load_classes(cfg["classes"]))) + # 单测里不下载 ImageNet 权重 + if model_cfg.get("backbone", {}).get("name") == "torchvision": + model_cfg["backbone"] = {**model_cfg["backbone"], "pretrained": False} + model = build_model(model_cfg) + with torch.no_grad(): + out = model(torch.randn(1, 3, 64, 64)) + assert "heatmap" in out + + simclr_backbone = cfg.get_path("simclr.backbone") + if simclr_backbone: + backbone_cfg = dict(simclr_backbone) + if backbone_cfg.get("name") == "torchvision": + backbone_cfg["pretrained"] = False + backbone = BACKBONES.build(backbone_cfg) + assert backbone.out_channels > 0 + + +def test_eval_config_matches_its_training_backbone(): + """评估配置构建出的骨架必须与训练配置一致，否则权重根本加载不上。""" + for train_cfg, eval_cfg in [ + ("configs/train/baseline.yaml", "configs/eval/baseline.yaml"), + ("configs/train/resnet18.yaml", "configs/eval/resnet18.yaml"), + ]: + assert load_config(train_cfg)["model"] == load_config(eval_cfg)["model"], eval_cfg diff --git a/utils/config.py b/utils/config.py index f573b5b..3f453e6 100644 --- a/utils/config.py +++ b/utils/config.py @@ -64,21 +64,25 @@ def deep_merge(base: Mapping, override: Mapping) -> dict: merged = dict(base) for key, value in override.items(): if isinstance(value, Mapping) and value.get(REPLACE_KEY): - merged[key] = _strip_replace_key(value) + merged[key] = dict(value) # 整块替换；标记本身留到最后统一剥掉 elif key in merged and isinstance(merged[key], Mapping) and isinstance(value, Mapping): merged[key] = deep_merge(merged[key], value) else: - merged[key] = _strip_replace_key(value) if isinstance(value, Mapping) else value + merged[key] = value return merged -def _strip_replace_key(value: Mapping) -> dict: - """去掉 `_replace_` 标记本身，避免它漏进组件的构造参数。""" - return { - k: (_strip_replace_key(v) if isinstance(v, Mapping) else v) - for k, v in value.items() - if k != REPLACE_KEY - } +def strip_replace_keys(value: Any) -> Any: + """剥掉所有 `_replace_` 标记，避免它漏进组件的构造参数。 + + 必须等**全部合并做完**再剥：`_load_raw` 对每个 base 都会先 `deep_merge({}, data)` + 一次，若在那时就剥掉，标记还没来得及参与真正的替换判断就没了。 + """ + if isinstance(value, Mapping): + return {k: strip_replace_keys(v) for k, v in value.items() if k != REPLACE_KEY} + if isinstance(value, list): + return [strip_replace_keys(v) for v in value] + return value def _load_raw(path: Path, seen: set[Path]) -> dict: @@ -132,5 +136,6 @@ def apply_overrides(cfg: dict, overrides: Iterable[str]) -> dict: def load_config(path: str | Path, overrides: Iterable[str] = ()) -> Config: cfg = _load_raw(Path(path), set()) cfg = apply_overrides(cfg, overrides) + cfg = strip_replace_keys(cfg) cfg.setdefault("config_path", str(Path(path).resolve())) return Config(cfg)