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>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
f4eff8beee
commit
3a36aad6c3
@@ -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",
|
||||
]
|
||||
@@ -0,0 +1,4 @@
|
||||
from models.backbones.resnet import TorchvisionBackbone
|
||||
from models.backbones.simple_conv import TinyConvNet
|
||||
|
||||
__all__ = ["TinyConvNet", "TorchvisionBackbone"]
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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))
|
||||
@@ -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
|
||||
@@ -0,0 +1,4 @@
|
||||
from models.heads.center_head import CenterHead
|
||||
from models.heads.classification import ClassificationHead
|
||||
|
||||
__all__ = ["CenterHead", "ClassificationHead"]
|
||||
@@ -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),
|
||||
}
|
||||
@@ -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))
|
||||
@@ -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))
|
||||
@@ -0,0 +1,3 @@
|
||||
from models.necks.upsample import UpsampleNeck
|
||||
|
||||
__all__ = ["UpsampleNeck"]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
Reference in New Issue
Block a user