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>
28 lines
925 B
Python
28 lines
925 B
Python
"""图像级缺陷分类头。
|
||
|
||
输出维度由 `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))
|