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
-155
@@ -1,9 +1,10 @@
|
||||
"""指标：分类与定位。"""
|
||||
"""训练过程中的通用统计。
|
||||
|
||||
分类任务的指标（混淆矩阵、macro-F1 等）在 engine/classification.py 里。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class AverageMeter:
|
||||
def __init__(self) -> None:
|
||||
@@ -17,155 +18,3 @@ class AverageMeter:
|
||||
@property
|
||||
def avg(self) -> float:
|
||||
return self.total / self.count if self.count else 0.0
|
||||
|
||||
|
||||
def box_iou(boxes_a: torch.Tensor, boxes_b: torch.Tensor) -> torch.Tensor:
|
||||
"""(N, 4) 与 (M, 4) 的两两 IoU，返回 (N, M)。"""
|
||||
if boxes_a.numel() == 0 or boxes_b.numel() == 0:
|
||||
return boxes_a.new_zeros((boxes_a.shape[0], boxes_b.shape[0]))
|
||||
area_a = (boxes_a[:, 2] - boxes_a[:, 0]).clamp(min=0) * (boxes_a[:, 3] - boxes_a[:, 1]).clamp(min=0)
|
||||
area_b = (boxes_b[:, 2] - boxes_b[:, 0]).clamp(min=0) * (boxes_b[:, 3] - boxes_b[:, 1]).clamp(min=0)
|
||||
|
||||
lt = torch.max(boxes_a[:, None, :2], boxes_b[None, :, :2])
|
||||
rb = torch.min(boxes_a[:, None, 2:], boxes_b[None, :, 2:])
|
||||
wh = (rb - lt).clamp(min=0)
|
||||
inter = wh[..., 0] * wh[..., 1]
|
||||
return inter / (area_a[:, None] + area_b[None, :] - inter).clamp(min=1e-6)
|
||||
|
||||
|
||||
class DetectionMetric:
|
||||
"""按 IoU 阈值做贪心匹配，统计 precision / recall / F1 与平均 IoU。
|
||||
|
||||
这是工程上够用的近似指标；正式对比模型时建议再跑一次 COCO mAP。
|
||||
"""
|
||||
|
||||
def __init__(self, iou_threshold: float = 0.5) -> None:
|
||||
self.iou_threshold = iou_threshold
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
self.tp = 0
|
||||
self.fp = 0
|
||||
self.fn = 0
|
||||
self.iou_sum = 0.0
|
||||
|
||||
def update(
|
||||
self,
|
||||
pred_boxes: torch.Tensor,
|
||||
pred_labels: torch.Tensor,
|
||||
gt_boxes: torch.Tensor,
|
||||
gt_labels: torch.Tensor,
|
||||
) -> None:
|
||||
if gt_boxes.numel() == 0:
|
||||
self.fp += int(pred_boxes.shape[0])
|
||||
return
|
||||
if pred_boxes.numel() == 0:
|
||||
self.fn += int(gt_boxes.shape[0])
|
||||
return
|
||||
|
||||
ious = box_iou(pred_boxes.float().cpu(), gt_boxes.float().cpu())
|
||||
same_class = pred_labels.cpu()[:, None] == gt_labels.cpu()[None, :]
|
||||
ious = ious * same_class.float()
|
||||
|
||||
matched_gt: set[int] = set()
|
||||
for pred_idx in range(ious.shape[0]):
|
||||
best_iou, best_gt = ious[pred_idx].max(dim=0)
|
||||
gt_idx = int(best_gt)
|
||||
if float(best_iou) >= self.iou_threshold and gt_idx not in matched_gt:
|
||||
matched_gt.add(gt_idx)
|
||||
self.tp += 1
|
||||
self.iou_sum += float(best_iou)
|
||||
else:
|
||||
self.fp += 1
|
||||
self.fn += int(gt_boxes.shape[0]) - len(matched_gt)
|
||||
|
||||
def compute(self) -> dict[str, float]:
|
||||
precision = self.tp / max(self.tp + self.fp, 1)
|
||||
recall = self.tp / max(self.tp + self.fn, 1)
|
||||
f1 = 2 * precision * recall / max(precision + recall, 1e-9)
|
||||
return {
|
||||
"precision": precision,
|
||||
"recall": recall,
|
||||
"f1": f1,
|
||||
"mean_iou": self.iou_sum / max(self.tp, 1),
|
||||
}
|
||||
|
||||
|
||||
class MultiLabelMetric:
|
||||
"""多类模式：逐类阈值化后统计 micro precision / recall / F1 与整图完全正确率。"""
|
||||
|
||||
def __init__(self, threshold: float = 0.5) -> None:
|
||||
self.threshold = threshold
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
self.tp = 0
|
||||
self.fp = 0
|
||||
self.fn = 0
|
||||
self.exact = 0
|
||||
self.total = 0
|
||||
|
||||
def update(self, logits: torch.Tensor, targets: torch.Tensor) -> None:
|
||||
pred = (logits.sigmoid() >= self.threshold).float().cpu()
|
||||
gt = targets.float().cpu()
|
||||
self.tp += int((pred * gt).sum())
|
||||
self.fp += int((pred * (1 - gt)).sum())
|
||||
self.fn += int(((1 - pred) * gt).sum())
|
||||
self.exact += int((pred == gt).all(dim=1).sum())
|
||||
self.total += int(gt.shape[0])
|
||||
|
||||
def compute(self) -> dict[str, float]:
|
||||
precision = self.tp / max(self.tp + self.fp, 1)
|
||||
recall = self.tp / max(self.tp + self.fn, 1)
|
||||
return {
|
||||
"cls_precision": precision,
|
||||
"cls_recall": recall,
|
||||
"cls_f1": 2 * precision * recall / max(precision + recall, 1e-9),
|
||||
"cls_exact_match": self.exact / max(self.total, 1),
|
||||
}
|
||||
|
||||
|
||||
class SingleLabelMetric:
|
||||
"""单类模式：互斥分类的准确率与 macro precision / recall / F1。
|
||||
|
||||
指标名与 `MultiLabelMetric` 保持一致（cls_precision / cls_recall / cls_f1），
|
||||
因此切换 label_mode 后 `train.monitor` 不需要改。
|
||||
"""
|
||||
|
||||
def __init__(self, num_outputs: int) -> None:
|
||||
self.num_outputs = num_outputs
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
self.tp = torch.zeros(self.num_outputs)
|
||||
self.fp = torch.zeros(self.num_outputs)
|
||||
self.fn = torch.zeros(self.num_outputs)
|
||||
self.correct = 0
|
||||
self.total = 0
|
||||
|
||||
def update(self, logits: torch.Tensor, targets: torch.Tensor) -> None:
|
||||
pred = logits.argmax(dim=1).cpu()
|
||||
gt = targets.long().cpu()
|
||||
for cls in range(self.num_outputs):
|
||||
pred_cls = pred == cls
|
||||
gt_cls = gt == cls
|
||||
self.tp[cls] += int((pred_cls & gt_cls).sum())
|
||||
self.fp[cls] += int((pred_cls & ~gt_cls).sum())
|
||||
self.fn[cls] += int((~pred_cls & gt_cls).sum())
|
||||
self.correct += int((pred == gt).sum())
|
||||
self.total += int(gt.numel())
|
||||
|
||||
def compute(self) -> dict[str, float]:
|
||||
# 只对验证集中出现过的类别做 macro 平均，避免缺席类别把指标拉到 0。
|
||||
present = (self.tp + self.fn) > 0
|
||||
precision = self.tp / (self.tp + self.fp).clamp(min=1)
|
||||
recall = self.tp / (self.tp + self.fn).clamp(min=1)
|
||||
f1 = 2 * precision * recall / (precision + recall).clamp(min=1e-9)
|
||||
if present.any():
|
||||
precision, recall, f1 = precision[present], recall[present], f1[present]
|
||||
return {
|
||||
"cls_accuracy": self.correct / max(self.total, 1),
|
||||
"cls_precision": float(precision.mean()),
|
||||
"cls_recall": float(recall.mean()),
|
||||
"cls_f1": float(f1.mean()),
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user