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:
co-authored by
Claude Opus 5
parent
f4eff8beee
commit
3a36aad6c3
@@ -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
|
||||
Reference in New Issue
Block a user