feat(models): 纯卷积缺陷分类定位模型

面向实时设备部署，全程只用 convolution，不含 attention / transformer。

- backbones: 自研 tiny_conv（深度可分离，默认）与 torchvision 封装（resnet/mobilenet/shufflenet）
- necks: nearest 上采样 + 3x3 卷积，把特征恢复到 stride 4，兼顾小缺陷定位与速度
- heads: anchor-free 中心点定位头（heatmap + wh + offset）与图像级分类头；
  分类头输出维度随 label_mode 切换（多类多热 / 单类互斥含背景位）
- postprocess: 3x3 max-pool 峰值抑制替代 NMS 循环，无循环、可导出 TensorRT / OpenVINO

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

models/__init__.py

@@ -0,0 +1,20 @@
+"""模型骨架：backbone -> neck -> heads。
+
+约束：部署在实时设备上，只使用基于 convolution 的结构（不含 attention / transformer）。
+"""
+
+from models import backbones, heads, necks  # noqa: F401  触发注册
+from models.build import build_model
+from models.defect_net import DefectNet
+from models.postprocess import decode_heatmap
+from models.registry import BACKBONES, HEADS, MODELS, NECKS
+
+__all__ = [
+    "build_model",
+    "DefectNet",
+    "decode_heatmap",
+    "BACKBONES",
+    "NECKS",
+    "HEADS",
+    "MODELS",
+]

models/backbones/__init__.py

@@ -0,0 +1,4 @@
+from models.backbones.resnet import TorchvisionBackbone
+from models.backbones.simple_conv import TinyConvNet
+
+__all__ = ["TinyConvNet", "TorchvisionBackbone"]

models/backbones/resnet.py

@@ -0,0 +1,73 @@
+"""torchvision 卷积骨架封装（resnet / mobilenet 等纯卷积模型）。"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from models.registry import BACKBONES
+
+# 仅允许纯卷积骨架；ViT/Swin 等不在部署范围内。
+_SUPPORTED = {
+    "resnet18": 512,
+    "resnet34": 512,
+    "resnet50": 2048,
+    "mobilenet_v3_small": 576,
+    "mobilenet_v3_large": 960,
+    "shufflenet_v2_x1_0": 1024,
+}
+
+
+@BACKBONES.register("torchvision")
+class TorchvisionBackbone(nn.Module):
+    """输出 stride=32 的单尺度特征。
+
+    Args:
+        arch: `_SUPPORTED` 中的模型名。
+        pretrained: 是否加载 ImageNet 权重（需要联网或本地缓存）。
+        in_channels: 输入通道，非 3 时会把第一层卷积权重按通道求均值复制。
+        frozen: 是否冻结骨架参数。
+    """
+
+    def __init__(
+        self,
+        arch: str = "resnet18",
+        pretrained: bool = True,
+        in_channels: int = 3,
+        frozen: bool = False,
+    ) -> None:
+        super().__init__()
+        if arch not in _SUPPORTED:
+            raise ValueError(f"不支持的骨架 {arch}，可选 {sorted(_SUPPORTED)}")
+
+        import torchvision
+
+        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
+
+        if in_channels != 3:
+            self._patch_first_conv(in_channels)
+
+        self.out_channels = _SUPPORTED[arch]
+        self.out_stride = 32
+
+        if frozen:
+            for param in self.parameters():
+                param.requires_grad_(False)
+
+    def _patch_first_conv(self, in_channels: int) -> None:
+        for module in self.features.modules():
+            if isinstance(module, nn.Conv2d):
+                weight = module.weight.data.mean(dim=1, keepdim=True).repeat(1, in_channels, 1, 1)
+                module.in_channels = in_channels
+                module.weight = nn.Parameter(weight * (3.0 / in_channels))
+                return
+        raise RuntimeError("骨架中未找到 Conv2d，无法调整输入通道")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.features(x)

models/backbones/simple_conv.py

@@ -0,0 +1,62 @@
+"""自研轻量卷积骨架，面向实时推理。"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+from torch import nn
+
+from models.layers import ConvBNAct, DepthwiseSeparableConv, ResidualBlock
+from models.registry import BACKBONES
+
+_BLOCKS = {"dw": DepthwiseSeparableConv, "res": ResidualBlock}
+
+
+@BACKBONES.register("tiny_conv")
+class TinyConvNet(nn.Module):
+    """4 个下采样 stage，输出 stride=16 的单尺度特征。
+
+    Args:
+        in_channels: 输入通道（薄膜灰度图为 1，RGB 为 3）。
+        widths: 每个 stage 的输出通道。
+        depths: 每个 stage 的 block 数量。
+        block: "dw"(深度可分离，更快) 或 "res"(残差，更强)。
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        widths: Sequence[int] = (32, 64, 128, 256),
+        depths: Sequence[int] = (1, 2, 2, 2),
+        block: str = "dw",
+    ) -> None:
+        super().__init__()
+        if len(widths) != len(depths):
+            raise ValueError("widths 与 depths 长度必须一致")
+        if block not in _BLOCKS:
+            raise ValueError(f"未知 block 类型: {block}，可选 {sorted(_BLOCKS)}")
+        block_cls = _BLOCKS[block]
+
+        self.stem = ConvBNAct(in_channels, widths[0], kernel_size=3, stride=2)
+
+        stages: list[nn.Module] = []
+        prev = widths[0]
+        for stage_idx, (width, depth) in enumerate(zip(widths, depths, strict=True)):
+            layers: list[nn.Module] = []
+            for block_idx in range(depth):
+                stride = 2 if (block_idx == 0 and stage_idx > 0) else 1
+                layers.append(block_cls(prev, width, stride))
+                prev = width
+            stages.append(nn.Sequential(*layers))
+        self.stages = nn.ModuleList(stages)
+
+        self.out_channels = prev
+        # stem 下采样 2 倍，其后每个 stage(除第一个)各下采样 2 倍。
+        self.out_stride = 2 * 2 ** (len(widths) - 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.stem(x)
+        for stage in self.stages:
+            x = stage(x)
+        return x

models/build.py

@@ -0,0 +1,15 @@
+"""按配置构建模型。"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from torch import nn
+
+from models.registry import MODELS
+
+
+def build_model(cfg: Mapping[str, Any]) -> nn.Module:
+    """`cfg` 至少包含 `name` 与 `num_classes`，其余透传给模型构造函数。"""
+    return MODELS.build(dict(cfg))

models/defect_net.py

@@ -0,0 +1,83 @@
+"""缺陷分类 + 定位主模型。"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+from torch import nn
+
+from models.heads.center_head import CenterHead
+from models.heads.classification import ClassificationHead
+from models.registry import BACKBONES, MODELS, NECKS
+from utils.labels import LABEL_MODES, MULTILABEL, check_label_mode
+
+__all__ = ["DefectNet", "classification_outputs", "LABEL_MODES"]
+
+
+def classification_outputs(num_classes: int, label_mode: str) -> int:
+    """分类头的输出维度；multiclass 需要额外一维表示 "无缺陷"。"""
+    check_label_mode(label_mode)
+    return num_classes if label_mode == MULTILABEL else num_classes + 1
+
+
+@MODELS.register("defect_net")
+class DefectNet(nn.Module):
+    """backbone -> neck -> {分类头, 中心点定位头}。
+
+    Args:
+        num_classes: 缺陷类别数（不含背景）。
+        backbone: backbone 注册表的配置，如 `{"name": "tiny_conv", ...}`。
+        neck: neck 配置，`None` 表示直接在骨架特征上出图。
+        head: CenterHead 的额外参数。
+        with_classification: 是否输出图像级分类 logits。
+        label_mode: 图像级标签形式，由 yaml 配置决定。
+            "multilabel" —— 多类：一张图可同时含多种缺陷，逐类 sigmoid + BCE；
+            "multiclass" —— 单类：一张图只归一个类别（含 "无缺陷"），softmax + CE。
+            仅影响分类头与分类损失；定位头始终按 num_classes 个类别输出 heatmap。
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        backbone: Mapping[str, Any] | None = None,
+        neck: Mapping[str, Any] | None = None,
+        head: Mapping[str, Any] | None = None,
+        with_classification: bool = True,
+        classification: Mapping[str, Any] | None = None,
+        label_mode: str = "multilabel",
+    ) -> None:
+        super().__init__()
+        self.num_classes = num_classes
+        self.label_mode = check_label_mode(label_mode)
+
+        self.backbone = BACKBONES.build(dict(backbone or {"name": "tiny_conv"}))
+        stride = self.backbone.out_stride
+        channels = self.backbone.out_channels
+
+        if neck:
+            self.neck: nn.Module = NECKS.build(dict(neck), in_channels=channels)
+            channels = self.neck.out_channels
+            stride = stride // self.neck.scale_factor
+        else:
+            self.neck = nn.Identity()
+
+        self.out_stride = stride
+        self.center_head = CenterHead(channels, num_classes, **dict(head or {}))
+        self.cls_head = (
+            ClassificationHead(
+                channels,
+                classification_outputs(num_classes, label_mode),
+                **dict(classification or {}),
+            )
+            if with_classification
+            else None
+        )
+
+    def forward(self, images: torch.Tensor) -> dict[str, torch.Tensor]:
+        feat = self.neck(self.backbone(images))
+        outputs = self.center_head(feat)
+        if self.cls_head is not None:
+            outputs["logits"] = self.cls_head(feat)
+        return outputs

models/heads/__init__.py

@@ -0,0 +1,4 @@
+from models.heads.center_head import CenterHead
+from models.heads.classification import ClassificationHead
+
+__all__ = ["CenterHead", "ClassificationHead"]

models/heads/center_head.py

@@ -0,0 +1,51 @@
+"""Anchor-free 中心点定位头（CenterNet 风格，纯卷积）。
+
+输出:
+    heatmap: (B, num_classes, H, W)  各类别缺陷中心的置信度（logits）
+    wh:      (B, 2, H, W)            中心点处的框宽高（特征图尺度）
+    offset:  (B, 2, H, W)            下采样量化误差补偿
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+
+from models.layers import ConvBNAct
+from models.registry import HEADS
+
+
+def _branch(in_channels: int, hidden: int, out_channels: int) -> nn.Sequential:
+    return nn.Sequential(
+        ConvBNAct(in_channels, hidden, kernel_size=3),
+        nn.Conv2d(hidden, out_channels, kernel_size=1),
+    )
+
+
+@HEADS.register("center")
+class CenterHead(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        num_classes: int,
+        hidden_channels: int = 64,
+        prior_prob: float = 0.01,
+    ) -> None:
+        super().__init__()
+        self.num_classes = num_classes
+        self.heatmap = _branch(in_channels, hidden_channels, num_classes)
+        self.wh = _branch(in_channels, hidden_channels, 2)
+        self.offset = _branch(in_channels, hidden_channels, 2)
+
+        # heatmap 正样本极稀疏，用先验概率初始化最后一层 bias，稳定早期训练。
+        bias = -math.log((1 - prior_prob) / prior_prob)
+        nn.init.constant_(self.heatmap[-1].bias, bias)
+
+    def forward(self, feat: torch.Tensor) -> dict[str, torch.Tensor]:
+        return {
+            "heatmap": self.heatmap(feat),
+            "wh": self.wh(feat),
+            "offset": self.offset(feat),
+        }

models/heads/classification.py

@@ -0,0 +1,27 @@
+"""图像级缺陷分类头。
+
+输出维度由 `label_mode` 决定，见 `models.defect_net.DefectNet`：
+    multilabel: num_classes      —— 每一维独立 sigmoid，全 0 表示 OK 片
+    multiclass: num_classes + 1  —— softmax 互斥，最后一维为 "无缺陷"
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from models.registry import HEADS
+
+
+@HEADS.register("classification")
+class ClassificationHead(nn.Module):
+    def __init__(self, in_channels: int, num_outputs: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        self.fc = nn.Linear(in_channels, num_outputs)
+        self.num_outputs = num_outputs
+
+    def forward(self, feat: torch.Tensor) -> torch.Tensor:
+        x = self.pool(feat).flatten(1)
+        return self.fc(self.drop(x))

models/layers.py

@@ -0,0 +1,51 @@
+"""基础卷积算子。全部为纯卷积结构，便于导出 ONNX / TensorRT。"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class ConvBNAct(nn.Sequential):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        groups: int = 1,
+        act: bool = True,
+    ) -> None:
+        padding = kernel_size // 2
+        layers: list[nn.Module] = [
+            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, groups=groups, bias=False),
+            nn.BatchNorm2d(out_channels),
+        ]
+        if act:
+            layers.append(nn.ReLU(inplace=True))
+        super().__init__(*layers)
+
+
+class DepthwiseSeparableConv(nn.Sequential):
+    """深度可分离卷积，实时设备上的默认选择。"""
+
+    def __init__(self, in_channels: int, out_channels: int, stride: int = 1) -> None:
+        super().__init__(
+            ConvBNAct(in_channels, in_channels, 3, stride, groups=in_channels),
+            ConvBNAct(in_channels, out_channels, 1),
+        )
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, stride: int = 1) -> None:
+        super().__init__()
+        self.conv1 = ConvBNAct(in_channels, out_channels, 3, stride)
+        self.conv2 = ConvBNAct(out_channels, out_channels, 3, 1, act=False)
+        self.act = nn.ReLU(inplace=True)
+        if stride != 1 or in_channels != out_channels:
+            self.shortcut: nn.Module = ConvBNAct(in_channels, out_channels, 1, stride, act=False)
+        else:
+            self.shortcut = nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(self.conv2(self.conv1(x)) + self.shortcut(x))

models/necks/__init__.py

@@ -0,0 +1,3 @@
+from models.necks.upsample import UpsampleNeck
+
+__all__ = ["UpsampleNeck"]

models/necks/upsample.py

@@ -0,0 +1,33 @@
+"""上采样 neck：把骨架的低分辨率特征恢复到定位所需的 stride。"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from models.layers import ConvBNAct
+from models.registry import NECKS
+
+
+@NECKS.register("upsample")
+class UpsampleNeck(nn.Module):
+    """`num_upsamples` 次 2 倍上采样，每次接一个 3x3 卷积平滑。
+
+    使用 nearest + conv 而非转置卷积：避免棋盘伪影，且在推理引擎上更快。
+    """
+
+    def __init__(self, in_channels: int, channels: int = 64, num_upsamples: int = 2) -> None:
+        super().__init__()
+        self.lateral = ConvBNAct(in_channels, channels, kernel_size=1)
+        self.blocks = nn.ModuleList(
+            [ConvBNAct(channels, channels, kernel_size=3) for _ in range(num_upsamples)]
+        )
+        self.out_channels = channels
+        self.scale_factor = 2**num_upsamples
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.lateral(x)
+        for block in self.blocks:
+            x = block(F.interpolate(x, scale_factor=2.0, mode="nearest"))
+        return x

models/postprocess.py

@@ -0,0 +1,68 @@
+"""推理后处理：从 heatmap 解码出缺陷框。"""
+
+from __future__ import annotations
+
+import torch
+from torch.nn import functional as F
+
+
+def heatmap_nms(heatmap: torch.Tensor, kernel: int = 3) -> torch.Tensor:
+    """用 max-pool 做峰值抑制，替代 NMS（无循环，可导出）。"""
+    pooled = F.max_pool2d(heatmap, kernel, stride=1, padding=kernel // 2)
+    return heatmap * (pooled == heatmap).float()
+
+
+@torch.no_grad()
+def decode_heatmap(
+    heatmap: torch.Tensor,
+    wh: torch.Tensor,
+    offset: torch.Tensor,
+    stride: int,
+    topk: int = 50,
+    score_threshold: float = 0.3,
+) -> list[dict[str, torch.Tensor]]:
+    """解码为原图坐标下的检测结果。
+
+    Args:
+        heatmap: (B, C, H, W) logits。
+        wh / offset: (B, 2, H, W)，单位为特征图像素。
+        stride: 特征图相对原图的下采样倍数。
+
+    Returns:
+        每张图一个 dict：`boxes` (N, 4) xyxy、`scores` (N,)、`labels` (N,)。
+    """
+    scores_map = heatmap_nms(heatmap.sigmoid())
+    batch, _num_classes, height, width = scores_map.shape
+
+    flat = scores_map.view(batch, -1)
+    topk = min(topk, flat.shape[1])
+    topk_scores, topk_idx = flat.topk(topk, dim=1)
+
+    labels = torch.div(topk_idx, height * width, rounding_mode="floor")
+    pixel_idx = topk_idx % (height * width)
+    ys = torch.div(pixel_idx, width, rounding_mode="floor").float()
+    xs = (pixel_idx % width).float()
+
+    wh_flat = wh.view(batch, 2, -1)
+    offset_flat = offset.view(batch, 2, -1)
+    gather_idx = pixel_idx.unsqueeze(1).expand(-1, 2, -1)
+    pred_wh = wh_flat.gather(2, gather_idx)
+    pred_offset = offset_flat.gather(2, gather_idx)
+
+    cx = (xs + pred_offset[:, 0]) * stride
+    cy = (ys + pred_offset[:, 1]) * stride
+    half_w = pred_wh[:, 0].clamp(min=0) * stride / 2
+    half_h = pred_wh[:, 1].clamp(min=0) * stride / 2
+    boxes = torch.stack([cx - half_w, cy - half_h, cx + half_w, cy + half_h], dim=-1)
+
+    results = []
+    for i in range(batch):
+        keep = topk_scores[i] >= score_threshold
+        results.append(
+            {
+                "boxes": boxes[i][keep],
+                "scores": topk_scores[i][keep],
+                "labels": labels[i][keep],
+            }
+        )
+    return results

models/registry.py

@@ -0,0 +1,10 @@
+"""模型组件注册表。"""
+
+from utils.registry import Registry
+
+BACKBONES = Registry("backbones")
+NECKS = Registry("necks")
+HEADS = Registry("heads")
+MODELS = Registry("models")
+
+__all__ = ["BACKBONES", "NECKS", "HEADS", "MODELS"]