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>
74 lines
2.3 KiB
Python
74 lines
2.3 KiB
Python
"""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)
|