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:
co-authored by
Claude Opus 5
parent
57446c2f10
commit
86c86b105c
@@ -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),
|
||||
}
|
||||
Reference in New Issue
Block a user