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,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))
|
||||
Reference in New Issue
Block a user