Shaun ZhouandClaude Opus 5
3a36aad6c3
面向实时设备部署，全程只用 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>
63 lines
2.1 KiB
Python
63 lines
2.1 KiB
Python
"""自研轻量卷积骨架，面向实时推理。"""
|
|
|
|
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
|