From 3a36aad6c31c1443fa6c99a55f8abcf9bb5769cf Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Wed, 2 Sep 2026 17:54:55 +0800 Subject: [PATCH] =?UTF-8?q?feat(models):=20=E7=BA=AF=E5=8D=B7=E7=A7=AF?= =?UTF-8?q?=E7=BC=BA=E9=99=B7=E5=88=86=E7=B1=BB=E5=AE=9A=E4=BD=8D=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 面向实时设备部署，全程只用 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 --- models/__init__.py | 20 ++++++++ models/backbones/__init__.py | 4 ++ models/backbones/resnet.py | 73 +++++++++++++++++++++++++++++ models/backbones/simple_conv.py | 62 ++++++++++++++++++++++++ models/build.py | 15 ++++++ models/defect_net.py | 83 +++++++++++++++++++++++++++++++++ models/heads/__init__.py | 4 ++ models/heads/center_head.py | 51 ++++++++++++++++++++ models/heads/classification.py | 27 +++++++++++ models/layers.py | 51 ++++++++++++++++++++ models/necks/__init__.py | 3 ++ models/necks/upsample.py | 33 +++++++++++++ models/postprocess.py | 68 +++++++++++++++++++++++++++ models/registry.py | 10 ++++ 14 files changed, 504 insertions(+) create mode 100644 models/__init__.py create mode 100644 models/backbones/__init__.py create mode 100644 models/backbones/resnet.py create mode 100644 models/backbones/simple_conv.py create mode 100644 models/build.py create mode 100644 models/defect_net.py create mode 100644 models/heads/__init__.py create mode 100644 models/heads/center_head.py create mode 100644 models/heads/classification.py create mode 100644 models/layers.py create mode 100644 models/necks/__init__.py create mode 100644 models/necks/upsample.py create mode 100644 models/postprocess.py create mode 100644 models/registry.py diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000..6cab961 --- /dev/null +++ b/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", +] diff --git a/models/backbones/__init__.py b/models/backbones/__init__.py new file mode 100644 index 0000000..9396333 --- /dev/null +++ b/models/backbones/__init__.py @@ -0,0 +1,4 @@ +from models.backbones.resnet import TorchvisionBackbone +from models.backbones.simple_conv import TinyConvNet + +__all__ = ["TinyConvNet", "TorchvisionBackbone"] diff --git a/models/backbones/resnet.py b/models/backbones/resnet.py new file mode 100644 index 0000000..ad980a2 --- /dev/null +++ b/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) diff --git a/models/backbones/simple_conv.py b/models/backbones/simple_conv.py new file mode 100644 index 0000000..d0d00c1 --- /dev/null +++ b/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 diff --git a/models/build.py b/models/build.py new file mode 100644 index 0000000..ec7a527 --- /dev/null +++ b/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)) diff --git a/models/defect_net.py b/models/defect_net.py new file mode 100644 index 0000000..869d9f4 --- /dev/null +++ b/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 diff --git a/models/heads/__init__.py b/models/heads/__init__.py new file mode 100644 index 0000000..c80cbec --- /dev/null +++ b/models/heads/__init__.py @@ -0,0 +1,4 @@ +from models.heads.center_head import CenterHead +from models.heads.classification import ClassificationHead + +__all__ = ["CenterHead", "ClassificationHead"] diff --git a/models/heads/center_head.py b/models/heads/center_head.py new file mode 100644 index 0000000..8c35613 --- /dev/null +++ b/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), + } diff --git a/models/heads/classification.py b/models/heads/classification.py new file mode 100644 index 0000000..435b963 --- /dev/null +++ b/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)) diff --git a/models/layers.py b/models/layers.py new file mode 100644 index 0000000..f09a99d --- /dev/null +++ b/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)) diff --git a/models/necks/__init__.py b/models/necks/__init__.py new file mode 100644 index 0000000..f64a426 --- /dev/null +++ b/models/necks/__init__.py @@ -0,0 +1,3 @@ +from models.necks.upsample import UpsampleNeck + +__all__ = ["UpsampleNeck"] diff --git a/models/necks/upsample.py b/models/necks/upsample.py new file mode 100644 index 0000000..3806bf6 --- /dev/null +++ b/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 diff --git a/models/postprocess.py b/models/postprocess.py new file mode 100644 index 0000000..91f3983 --- /dev/null +++ b/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 diff --git a/models/registry.py b/models/registry.py new file mode 100644 index 0000000..6b968d0 --- /dev/null +++ b/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"]