From e581d9b3d7ee3a605801cb60462ae0870f98a50b Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Fri, 4 Sep 2026 15:28:49 +0800 Subject: [PATCH] =?UTF-8?q?fix(backbones):=20=E4=BF=AE=E5=A4=8D=20shufflen?= =?UTF-8?q?et=20=E6=97=A0=E6=B3=95=E6=9E=84=E5=BB=BA=EF=BC=8C=E5=B9=B6?= =?UTF-8?q?=E8=AE=A9=E6=8D=A2=E9=AA=A8=E6=9E=B6=E4=B8=8D=E5=86=8D=E8=B8=A9?= =?UTF-8?q?=E9=85=8D=E7=BD=AE=E5=90=88=E5=B9=B6=E7=9A=84=E5=9D=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - shufflenet_v2_x1_0 列在支持清单里却直接 AttributeError：ShuffleNetV2 既不是 扁平 children 也没有 .features，需要按 conv1/maxpool/stage2..4/conv5 自己拼。 该 bug 同时影响 SimCLR 与检测模型 - 新增参数化测试覆盖清单里的每个 arch：能前向、声明的 out_channels 与实际一致 - 配置支持 `_replace_: true`：子字典整块替换而不与 base 合并。换骨架时 tiny_conv 的 widths/depths 会残留并传给新骨架，此前无论用 --set 还是 _base_ 都绕不开，只能改原文件 - Registry.build 提前校验字段，报错直接指出是哪些字段不被接受、该怎么办， 取代裸 TypeError - README 补充可选骨架表与三种改法（含不该用 --set 改 name 的说明） Co-Authored-By: Claude Opus 5 --- README.md | 45 +++++++++++++++++++++++++++ docs/unsupervised-clustering.md | 1 + models/backbones/resnet.py | 24 +++++++++++--- tests/test_config.py | 37 ++++++++++++++++++++++ tests/test_models.py | 55 +++++++++++++++++++++++++++++++++ utils/config.py | 24 ++++++++++++-- utils/registry.py | 26 +++++++++++++++- 7 files changed, 204 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index bea3404..d7bf58a 100644 --- a/README.md +++ b/README.md @@ -353,6 +353,51 @@ view_transforms: # 生成两个视图的随机增强 > `RandomResizedCrop` 的 `scale` 下限别低于 0.4：裁太狠会让两个视图都只剩背景， > "背景 vs 背景"被当成正样本对，学出来的特征会退化。 +#### 换一个骨架 + +`simclr.backbone` 用的是和检测模型同一套注册表，可选： + +| `name` | `arch` | out_channels | 说明 | +| --- | --- | --- | --- | +| `tiny_conv` | — | 256（可调） | 自研深度可分离卷积，默认，最轻 | +| `torchvision` | `resnet18` / `resnet34` | 512 | 通用基线 | +| `torchvision` | `resnet50` | 2048 | 最强也最重，MPS 上注意 batch | +| `torchvision` | `mobilenet_v3_small` / `_large` | 576 / 960 | 面向实时部署 | +| `torchvision` | `shufflenet_v2_x1_0` | 1024 | 同上 | + +全是纯卷积（ViT/Swin 不在支持范围，本项目要部署在实时设备上）。改法是**整块替换** +`simclr.backbone`，因为两种骨架的字段完全不同（`widths/depths/block` vs `arch/pretrained`）： + +```yaml +# configs/cluster/film.yaml +simclr: + backbone: + name: torchvision + arch: resnet18 + pretrained: true # 先用 ImageNet 权重起步，再在本批数据上自监督微调 + in_channels: 3 +``` + +如果不想改原文件，用 `_base_` 继承时要加 `_replace_: true`——子字典默认是**合并**而不是 +替换，不加的话 base 里 `tiny_conv` 的 `widths/depths` 会残留下来传给新骨架： + +```yaml +_base_: [configs/cluster/film.yaml] +simclr: + backbone: + _replace_: true # 整块换掉，不与 base 合并 + name: torchvision + arch: resnet18 + pretrained: true +``` + +> **不要用 `--set simclr.backbone.name=torchvision`**：它只改名字，旧骨架的专有字段仍在， +> 会报 `构建 'torchvision' 时收到它不接受的字段 ['block', 'depths', 'widths']`。 +> 换组件请改配置文件。 + +注意 `simclr.feature_stages`：`torchvision` 系列的 `features` 是 `nn.Sequential`，可以填 +`[5, 6]` 这类中间层下标；`tiny_conv` 不是，必须保持 `null`（取最终输出）。 + SimCLR 训练同样接了 SwanLab（配置在 `simclr.swanlab`，字段与 `train.swanlab` 相同， 默认 `mode: local` 不需要登录）。训练跑起来后另开一个终端： diff --git a/docs/unsupervised-clustering.md b/docs/unsupervised-clustering.md index 45d60d9..a5abf11 100644 --- a/docs/unsupervised-clustering.md +++ b/docs/unsupervised-clustering.md @@ -743,6 +743,7 @@ $k=4$ 稳定性略低于 ImageNet（0.83 vs 0.98），符合"数据的主结构 | NT-Xent 损失 | `engine/simclr.py::nt_xent_loss` | | SimCLR 训练循环 | `engine/simclr.py::SimCLRTrainer` | | 断点续训（完整状态存取） | `engine/simclr.py::save_state` / `load_state` | +| 配置整块替换 `_replace_` | `utils/config.py::deep_merge` | | 投影头 $g_\varphi$ | `models/heads/projection.py::ProjectionHead` | | 设备解析（cuda/mps/cpu，auto 优先级） | `utils/device.py::resolve_device` | | AMP 随后端的默认值 | `utils/device.py::resolve_amp` / `default_amp` | diff --git a/models/backbones/resnet.py b/models/backbones/resnet.py index ad980a2..105371f 100644 --- a/models/backbones/resnet.py +++ b/models/backbones/resnet.py @@ -45,10 +45,7 @@ class TorchvisionBackbone(nn.Module): weights = "DEFAULT" if pretrained else None net = getattr(torchvision.models, arch)(weights=weights) - if arch.startswith("resnet"): - self.features = nn.Sequential(*list(net.children())[:-2]) - else: - self.features = net.features + self.features = self._extract_features(net, arch) if in_channels != 3: self._patch_first_conv(in_channels) @@ -60,6 +57,25 @@ class TorchvisionBackbone(nn.Module): for param in self.parameters(): param.requires_grad_(False) + @staticmethod + def _extract_features(net: nn.Module, arch: str) -> nn.Sequential: + """取到最后一个卷积 stage 为止，丢掉分类头。 + + 不同系列的 torchvision 模型结构不一样：resnet 是扁平的 children， + mobilenet 有现成的 `.features`，而 ShuffleNetV2 两者都不是——它把层 + 散放在 conv1/maxpool/stage2..4/conv5 上，得按顺序自己拼。 + """ + if arch.startswith("resnet"): + return nn.Sequential(*list(net.children())[:-2]) + if arch.startswith("shufflenet"): + return nn.Sequential( + net.conv1, net.maxpool, net.stage2, net.stage3, net.stage4, net.conv5 + ) + features = getattr(net, "features", None) + if features is None: + raise ValueError(f"不知道如何从 {arch} 取特征，请在 _extract_features 里补一条分支") + return features + def _patch_first_conv(self, in_channels: int) -> None: for module in self.features.modules(): if isinstance(module, nn.Conv2d): diff --git a/tests/test_config.py b/tests/test_config.py index 93d5a46..24525f4 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -46,3 +46,40 @@ def test_project_configs_are_loadable(): assert cfg["model"]["name"] == "defect_net" assert cfg["data"]["train"]["dataset"]["name"] == "film_defect" assert len(cfg["classes"]) > 0 + + +def test_replace_key_swaps_a_whole_subdict(tmp_path): + """换组件时要能整块替换，否则旧组件的专有字段会残留下来传给新组件。""" + (tmp_path / "base.yaml").write_text( + "backbone:\n name: tiny_conv\n widths: [32, 64]\n depths: [1, 2]\n block: dw\n", + encoding="utf-8", + ) + (tmp_path / "child.yaml").write_text( + "_base_: ./base.yaml\nbackbone:\n _replace_: true\n name: torchvision\n arch: resnet18\n", + encoding="utf-8", + ) + + cfg = load_config(tmp_path / "child.yaml") + + assert cfg["backbone"] == {"name": "torchvision", "arch": "resnet18"} + assert "_replace_" not in cfg["backbone"] # 标记本身不能漏进构造参数 + + +def test_without_replace_key_subdicts_still_merge(tmp_path): + (tmp_path / "base.yaml").write_text("backbone:\n name: tiny_conv\n block: dw\n", encoding="utf-8") + (tmp_path / "child.yaml").write_text( + "_base_: ./base.yaml\nbackbone:\n block: res\n", encoding="utf-8" + ) + + cfg = load_config(tmp_path / "child.yaml") + assert cfg["backbone"] == {"name": "tiny_conv", "block": "res"} + + +def test_registry_reports_incompatible_fields_actionably(): + """--set 只改 name 时旧字段会残留，报错要说清楚是哪些字段、该怎么办。""" + import pytest + + from models.registry import BACKBONES + + with pytest.raises(TypeError, match=r"不接受的字段.*widths"): + BACKBONES.build({"name": "torchvision", "arch": "resnet18", "widths": [32]}) diff --git a/tests/test_models.py b/tests/test_models.py index 26657cb..3123050 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -3,6 +3,7 @@ import pytest torch = pytest.importorskip("torch") from models import build_model # noqa: E402 +from models.backbones.resnet import _SUPPORTED as _SUPPORTED_ARCHS # noqa: E402 from models.postprocess import decode_heatmap # noqa: E402 @@ -90,3 +91,57 @@ def test_losses_decrease_on_a_single_batch(model_cfg): last = total assert last < first + + +@pytest.mark.parametrize("arch", sorted(_SUPPORTED_ARCHS)) +def test_every_supported_torchvision_arch_actually_works(arch): + """清单里列出的每个骨架都必须真能前向，且声明的 out_channels 与实际一致。 + + 回归测试：shufflenet_v2_x1_0 曾被列为支持，但它没有 `.features` 属性， + 构造时直接 AttributeError。 + """ + from models.registry import BACKBONES + + backbone = BACKBONES.build({"name": "torchvision", "arch": arch, "pretrained": False}).eval() + with torch.no_grad(): + feat = backbone(torch.randn(1, 3, 64, 64)) + + assert feat.dim() == 4 + assert feat.shape[1] == backbone.out_channels == _SUPPORTED_ARCHS[arch] + assert backbone.out_stride == 32 + + +def test_unsupported_arch_is_rejected(): + from models.registry import BACKBONES + + with pytest.raises(ValueError, match="不支持的骨架"): + BACKBONES.build({"name": "torchvision", "arch": "vit_b_16"}) + + +@pytest.mark.parametrize("arch", ["resnet18", "shufflenet_v2_x1_0", "mobilenet_v3_small"]) +def test_alternative_backbones_drive_simclr(arch): + """换骨架后 SimCLR 的投影头要能接上（out_channels 必须准）。""" + from models.heads.projection import ProjectionHead + from models.registry import BACKBONES + + backbone = BACKBONES.build({"name": "torchvision", "arch": arch, "pretrained": False}).eval() + head = ProjectionHead(backbone.out_channels, 32, 16) + with torch.no_grad(): + assert head(backbone(torch.randn(2, 3, 64, 64))).shape == (2, 16) + + +def test_alternative_backbone_drives_defect_net(): + """检测模型换骨架同样要能跑通（neck 的 in_channels 由 backbone 决定）。""" + model = build_model( + { + "name": "defect_net", + "num_classes": 3, + "backbone": {"name": "torchvision", "arch": "shufflenet_v2_x1_0", "pretrained": False}, + "neck": {"name": "upsample", "channels": 16, "num_upsamples": 3}, + "head": {"hidden_channels": 16}, + } + ).eval() + assert model.out_stride == 4 + with torch.no_grad(): + out = model(torch.randn(1, 3, 64, 64)) + assert out["heatmap"].shape == (1, 3, 16, 16) diff --git a/utils/config.py b/utils/config.py index 1e619b3..f573b5b 100644 --- a/utils/config.py +++ b/utils/config.py @@ -10,6 +10,7 @@ from typing import Any import yaml BASE_KEY = "_base_" +REPLACE_KEY = "_replace_" class Config(dict): @@ -53,16 +54,33 @@ class Config(dict): def deep_merge(base: Mapping, override: Mapping) -> dict: - """递归合并；override 中的标量/列表直接覆盖 base。""" + """递归合并；override 中的标量/列表直接覆盖 base。 + + 子字典默认是**合并**而不是替换，这在换组件时会出问题：把 backbone 从 + `tiny_conv` 换成 `torchvision`，base 里 tiny_conv 专有的 widths/depths + 会残留下来，传给新组件直接报错。这时在该字典里写 `_replace_: true`， + 表示整块替换掉 base 的同名字典。 + """ merged = dict(base) for key, value in override.items(): - if key in merged and isinstance(merged[key], Mapping) and isinstance(value, Mapping): + if isinstance(value, Mapping) and value.get(REPLACE_KEY): + merged[key] = _strip_replace_key(value) + elif key in merged and isinstance(merged[key], Mapping) and isinstance(value, Mapping): merged[key] = deep_merge(merged[key], value) else: - merged[key] = value + merged[key] = _strip_replace_key(value) if isinstance(value, Mapping) else 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 _load_raw(path: Path, seen: set[Path]) -> dict: path = path.resolve() if path in seen: diff --git a/utils/registry.py b/utils/registry.py index 62a25ed..6b0a85a 100644 --- a/utils/registry.py +++ b/utils/registry.py @@ -2,6 +2,7 @@ from __future__ import annotations +import inspect from collections.abc import Callable, Iterator from typing import Any, TypeVar @@ -38,7 +39,30 @@ class Registry: name = cfg.pop("name", None) if name is None: raise KeyError(f"构建 {self.name} 需要 'name' 字段: {cfg}") - return self.get(name)(**{**cfg, **kwargs}) + target = self.get(name) + params = {**cfg, **kwargs} + self._check_params(name, target, params) + return target(**params) + + @staticmethod + def _check_params(name: str, target: Callable[..., Any], params: dict) -> None: + """提前挡下不认识的字段，给出比裸 TypeError 更可操作的提示。 + + 最常见的触发方式：用 `--set xxx.name=其他骨架` 只改了名字， + yaml 里原骨架的专有字段（如 tiny_conv 的 widths/depths）仍然留在配置里。 + """ + signature = inspect.signature(target) + if any(p.kind is p.VAR_KEYWORD for p in signature.parameters.values()): + return + unknown = sorted(set(params) - set(signature.parameters)) + if not unknown: + return + accepted = sorted(p for p in signature.parameters if p != "self") + raise TypeError( + f"构建 '{name}' 时收到它不接受的字段 {unknown}。它接受的是 {accepted}。" + "如果你是用 --set 改了 name，注意 yaml 里原来那个组件的专有字段仍会保留——" + "换组件请直接编辑配置文件、整块替换。" + ) def __contains__(self, name: object) -> bool: return name in self._entries