refactor: 按实际的五步工作流裁掉无关代码
实际工作流是：simclr 训练 -> 聚类 -> 人工归类 -> 分类训练 -> GradCAM/SHAP 生成
热力图与 box。据此移除不参与其中的检测/定位链路，共 31 个文件、约 2600 行：
- 模型：defect_net、postprocess、necks/、heads/center_head
- 数据：film_defect（YOLO 读取端）、target_encoder
- 引擎：trainer、evaluator、losses
- 入口：train.py、eval.py
- 脚本：convert_annotations、split_dataset、export_onnx、profile_inference
- 配置：configs/{train,eval,model,dataset}
- utils/labels（多类/单类标签模式，仅检测分类头使用）
共享模块按实际用量瘦身：engine/metrics 只留 AverageMeter（检测指标随之移除）；
datasets/yolo 只留写出侧（IMAGE_SUFFIXES 与 write_data_yaml）——第 5 步产出的
YOLO 标签是给外部检测训练用的交换格式，读取端不再需要。
裁剪过程中修掉一个真 bug：generate_boxes 的 --stage 默认 6（layer3）对 tiny_conv
这类非 nn.Sequential 骨架直接抛异常。现在自动退回骨架最终输出并警告，
另加 --stage -1 显式选择最终输出。
五步工作流已端到端实跑验证（simclr 1 轮 -> 聚类 -> 导出 -> 分类 1 轮 ->
gradcam 与 shap 两种方法各生成一遍）。
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
57446c2f10
commit
86c86b105c
+4
-9
@@ -1,22 +1,17 @@
|
||||
"""模型骨架：backbone -> neck -> heads。
|
||||
"""模型骨架：backbone + head。
|
||||
|
||||
约束：部署在实时设备上，只使用基于 convolution 的结构（不含 attention / transformer）。
|
||||
按项目约定只用 convolution，不含 attention / transformer。
|
||||
"""
|
||||
|
||||
from models import backbones, heads, necks # noqa: F401 触发注册
|
||||
from models import backbones, heads # noqa: F401 触发注册
|
||||
from models.build import build_model
|
||||
from models.classifier import DefectClassifier
|
||||
from models.defect_net import DefectNet
|
||||
from models.postprocess import decode_heatmap
|
||||
from models.registry import BACKBONES, HEADS, MODELS, NECKS
|
||||
from models.registry import BACKBONES, HEADS, MODELS
|
||||
|
||||
__all__ = [
|
||||
"build_model",
|
||||
"DefectNet",
|
||||
"DefectClassifier",
|
||||
"decode_heatmap",
|
||||
"BACKBONES",
|
||||
"NECKS",
|
||||
"HEADS",
|
||||
"MODELS",
|
||||
]
|
||||
|
||||
@@ -1,83 +0,0 @@
|
||||
"""缺陷分类 + 定位主模型。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from models.heads.center_head import CenterHead
|
||||
from models.heads.classification import ClassificationHead
|
||||
from models.registry import BACKBONES, MODELS, NECKS
|
||||
from utils.labels import LABEL_MODES, MULTILABEL, check_label_mode
|
||||
|
||||
__all__ = ["DefectNet", "classification_outputs", "LABEL_MODES"]
|
||||
|
||||
|
||||
def classification_outputs(num_classes: int, label_mode: str) -> int:
|
||||
"""分类头的输出维度；multiclass 需要额外一维表示 "无缺陷"。"""
|
||||
check_label_mode(label_mode)
|
||||
return num_classes if label_mode == MULTILABEL else num_classes + 1
|
||||
|
||||
|
||||
@MODELS.register("defect_net")
|
||||
class DefectNet(nn.Module):
|
||||
"""backbone -> neck -> {分类头, 中心点定位头}。
|
||||
|
||||
Args:
|
||||
num_classes: 缺陷类别数（不含背景）。
|
||||
backbone: backbone 注册表的配置，如 `{"name": "tiny_conv", ...}`。
|
||||
neck: neck 配置，`None` 表示直接在骨架特征上出图。
|
||||
head: CenterHead 的额外参数。
|
||||
with_classification: 是否输出图像级分类 logits。
|
||||
label_mode: 图像级标签形式，由 yaml 配置决定。
|
||||
"multilabel" —— 多类：一张图可同时含多种缺陷，逐类 sigmoid + BCE；
|
||||
"multiclass" —— 单类：一张图只归一个类别（含 "无缺陷"），softmax + CE。
|
||||
仅影响分类头与分类损失；定位头始终按 num_classes 个类别输出 heatmap。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_classes: int,
|
||||
backbone: Mapping[str, Any] | None = None,
|
||||
neck: Mapping[str, Any] | None = None,
|
||||
head: Mapping[str, Any] | None = None,
|
||||
with_classification: bool = True,
|
||||
classification: Mapping[str, Any] | None = None,
|
||||
label_mode: str = "multilabel",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_classes = num_classes
|
||||
self.label_mode = check_label_mode(label_mode)
|
||||
|
||||
self.backbone = BACKBONES.build(dict(backbone or {"name": "tiny_conv"}))
|
||||
stride = self.backbone.out_stride
|
||||
channels = self.backbone.out_channels
|
||||
|
||||
if neck:
|
||||
self.neck: nn.Module = NECKS.build(dict(neck), in_channels=channels)
|
||||
channels = self.neck.out_channels
|
||||
stride = stride // self.neck.scale_factor
|
||||
else:
|
||||
self.neck = nn.Identity()
|
||||
|
||||
self.out_stride = stride
|
||||
self.center_head = CenterHead(channels, num_classes, **dict(head or {}))
|
||||
self.cls_head = (
|
||||
ClassificationHead(
|
||||
channels,
|
||||
classification_outputs(num_classes, label_mode),
|
||||
**dict(classification or {}),
|
||||
)
|
||||
if with_classification
|
||||
else None
|
||||
)
|
||||
|
||||
def forward(self, images: torch.Tensor) -> dict[str, torch.Tensor]:
|
||||
feat = self.neck(self.backbone(images))
|
||||
outputs = self.center_head(feat)
|
||||
if self.cls_head is not None:
|
||||
outputs["logits"] = self.cls_head(feat)
|
||||
return outputs
|
||||
@@ -1,4 +1,4 @@
|
||||
from models.heads.center_head import CenterHead
|
||||
from models.heads.classification import ClassificationHead
|
||||
from models.heads.projection import ProjectionHead
|
||||
|
||||
__all__ = ["CenterHead", "ClassificationHead"]
|
||||
__all__ = ["ClassificationHead", "ProjectionHead"]
|
||||
|
||||
@@ -1,51 +0,0 @@
|
||||
"""Anchor-free 中心点定位头（CenterNet 风格，纯卷积）。
|
||||
|
||||
输出:
|
||||
heatmap: (B, num_classes, H, W) 各类别缺陷中心的置信度（logits）
|
||||
wh: (B, 2, H, W) 中心点处的框宽高（特征图尺度）
|
||||
offset: (B, 2, H, W) 下采样量化误差补偿
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from models.layers import ConvBNAct
|
||||
from models.registry import HEADS
|
||||
|
||||
|
||||
def _branch(in_channels: int, hidden: int, out_channels: int) -> nn.Sequential:
|
||||
return nn.Sequential(
|
||||
ConvBNAct(in_channels, hidden, kernel_size=3),
|
||||
nn.Conv2d(hidden, out_channels, kernel_size=1),
|
||||
)
|
||||
|
||||
|
||||
@HEADS.register("center")
|
||||
class CenterHead(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
num_classes: int,
|
||||
hidden_channels: int = 64,
|
||||
prior_prob: float = 0.01,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_classes = num_classes
|
||||
self.heatmap = _branch(in_channels, hidden_channels, num_classes)
|
||||
self.wh = _branch(in_channels, hidden_channels, 2)
|
||||
self.offset = _branch(in_channels, hidden_channels, 2)
|
||||
|
||||
# heatmap 正样本极稀疏，用先验概率初始化最后一层 bias，稳定早期训练。
|
||||
bias = -math.log((1 - prior_prob) / prior_prob)
|
||||
nn.init.constant_(self.heatmap[-1].bias, bias)
|
||||
|
||||
def forward(self, feat: torch.Tensor) -> dict[str, torch.Tensor]:
|
||||
return {
|
||||
"heatmap": self.heatmap(feat),
|
||||
"wh": self.wh(feat),
|
||||
"offset": self.offset(feat),
|
||||
}
|
||||
@@ -1,3 +0,0 @@
|
||||
from models.necks.upsample import UpsampleNeck
|
||||
|
||||
__all__ = ["UpsampleNeck"]
|
||||
@@ -1,33 +0,0 @@
|
||||
"""上采样 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
|
||||
@@ -1,68 +0,0 @@
|
||||
"""推理后处理：从 heatmap 解码出缺陷框。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def heatmap_nms(heatmap: torch.Tensor, kernel: int = 3) -> torch.Tensor:
|
||||
"""用 max-pool 做峰值抑制，替代 NMS（无循环，可导出）。"""
|
||||
pooled = F.max_pool2d(heatmap, kernel, stride=1, padding=kernel // 2)
|
||||
return heatmap * (pooled == heatmap).float()
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def decode_heatmap(
|
||||
heatmap: torch.Tensor,
|
||||
wh: torch.Tensor,
|
||||
offset: torch.Tensor,
|
||||
stride: int,
|
||||
topk: int = 50,
|
||||
score_threshold: float = 0.3,
|
||||
) -> list[dict[str, torch.Tensor]]:
|
||||
"""解码为原图坐标下的检测结果。
|
||||
|
||||
Args:
|
||||
heatmap: (B, C, H, W) logits。
|
||||
wh / offset: (B, 2, H, W)，单位为特征图像素。
|
||||
stride: 特征图相对原图的下采样倍数。
|
||||
|
||||
Returns:
|
||||
每张图一个 dict：`boxes` (N, 4) xyxy、`scores` (N,)、`labels` (N,)。
|
||||
"""
|
||||
scores_map = heatmap_nms(heatmap.sigmoid())
|
||||
batch, _num_classes, height, width = scores_map.shape
|
||||
|
||||
flat = scores_map.view(batch, -1)
|
||||
topk = min(topk, flat.shape[1])
|
||||
topk_scores, topk_idx = flat.topk(topk, dim=1)
|
||||
|
||||
labels = torch.div(topk_idx, height * width, rounding_mode="floor")
|
||||
pixel_idx = topk_idx % (height * width)
|
||||
ys = torch.div(pixel_idx, width, rounding_mode="floor").float()
|
||||
xs = (pixel_idx % width).float()
|
||||
|
||||
wh_flat = wh.view(batch, 2, -1)
|
||||
offset_flat = offset.view(batch, 2, -1)
|
||||
gather_idx = pixel_idx.unsqueeze(1).expand(-1, 2, -1)
|
||||
pred_wh = wh_flat.gather(2, gather_idx)
|
||||
pred_offset = offset_flat.gather(2, gather_idx)
|
||||
|
||||
cx = (xs + pred_offset[:, 0]) * stride
|
||||
cy = (ys + pred_offset[:, 1]) * stride
|
||||
half_w = pred_wh[:, 0].clamp(min=0) * stride / 2
|
||||
half_h = pred_wh[:, 1].clamp(min=0) * stride / 2
|
||||
boxes = torch.stack([cx - half_w, cy - half_h, cx + half_w, cy + half_h], dim=-1)
|
||||
|
||||
results = []
|
||||
for i in range(batch):
|
||||
keep = topk_scores[i] >= score_threshold
|
||||
results.append(
|
||||
{
|
||||
"boxes": boxes[i][keep],
|
||||
"scores": topk_scores[i][keep],
|
||||
"labels": labels[i][keep],
|
||||
}
|
||||
)
|
||||
return results
|
||||
+1
-2
@@ -3,8 +3,7 @@
|
||||
from utils.registry import Registry
|
||||
|
||||
BACKBONES = Registry("backbones")
|
||||
NECKS = Registry("necks")
|
||||
HEADS = Registry("heads")
|
||||
MODELS = Registry("models")
|
||||
|
||||
__all__ = ["BACKBONES", "NECKS", "HEADS", "MODELS"]
|
||||
__all__ = ["BACKBONES", "HEADS", "MODELS"]
|
||||
|
||||
Reference in New Issue
Block a user