feat: 新增缺陷分类训练与评估链路
在聚类 + 人工归类得到的数据集(data/<集>/<类别名>/*.png)上训练分类模型。 - datasets/classification.py:按目录组织的分类数据集,**逐类分层**切分 train/val (others 只有 65 张,整体随机切很容易让它在验证集里挂零),固定种子可复现 - models/classifier.py:卷积骨架 + 分类头,与 DefectNet 共用 backbone 注册表, 因此自监督产出的 simclr_backbone.pth 可用 init_from 直接初始化; frozen_backbone 可做线性探针 - engine/classification.py:训练循环 + 指标。主指标用 macro-F1 而非准确率—— 数据从 normal 2679 张到 others 65 张差 41 倍,全预测成 normal 也有 24% 准确率; 损失按类别频率加权,避免小类被淹没;输出逐类 P/R/F1 与混淆矩阵 - classify.py:训练/评估入口,--eval-only 可单独评估并导出 eval_report.json - configs/classify/film.yaml:预处理沿用聚类阶段的结论(拉伸 + 逐图标准化)。 抹掉亮度在这里尤其重要:有几类几乎只来自单一原始数据集,不抹掉的话模型可能 靠"认出这是哪个数据集"作弊 - 续训语义与检测/自监督一致:检查点只提供进度,超参以配置为准 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
33d28dcc7c
commit
33fdf396ee
@@ -0,0 +1,82 @@
|
||||
"""缺陷分类模型:卷积骨架 + 分类头。
|
||||
|
||||
与 `DefectNet` 共用同一套 backbone 注册表,所以自监督预训练出来的骨架
|
||||
(`runs/*/simclr_backbone.pth`)可以直接拿来初始化——那正是前面聚类阶段的产物。
|
||||
按项目约定只用 convolution,不含 attention / transformer。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from models.heads.classification import ClassificationHead
|
||||
from models.registry import BACKBONES, MODELS
|
||||
from utils.logging import get_logger
|
||||
|
||||
logger = get_logger("models.classifier")
|
||||
|
||||
|
||||
@MODELS.register("classifier")
|
||||
class DefectClassifier(nn.Module):
|
||||
"""Args:
|
||||
num_classes: 类别数。
|
||||
backbone: backbone 注册表的配置。
|
||||
classification: 分类头参数(`hidden`/`dropout` 等)。
|
||||
init_from: 自监督骨架权重路径;只加载 backbone,分类头仍随机初始化。
|
||||
frozen_backbone: 冻结骨架只训分类头(线性探针,用来单独衡量自监督特征的好坏)。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_classes: int,
|
||||
backbone: Mapping[str, Any] | None = None,
|
||||
classification: Mapping[str, Any] | None = None,
|
||||
init_from: str | Path | None = None,
|
||||
frozen_backbone: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_classes = num_classes
|
||||
self.backbone = BACKBONES.build(dict(backbone or {"name": "tiny_conv"}))
|
||||
self.head = ClassificationHead(
|
||||
self.backbone.out_channels, num_classes, **dict(classification or {})
|
||||
)
|
||||
|
||||
if init_from:
|
||||
self.load_backbone(init_from)
|
||||
if frozen_backbone:
|
||||
for param in self.backbone.parameters():
|
||||
param.requires_grad_(False)
|
||||
logger.info("骨架已冻结,只训练分类头(线性探针)")
|
||||
|
||||
def load_backbone(self, path: str | Path) -> None:
|
||||
"""加载自监督预训练的骨架权重。
|
||||
|
||||
只取 backbone 那部分:simclr_backbone.pth 里存的就是骨架 state_dict,
|
||||
分类头是新的、必须从头学。
|
||||
"""
|
||||
path = Path(path)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"预训练骨架不存在: {path}")
|
||||
state = torch.load(path, map_location="cpu", weights_only=False)
|
||||
weights = state.get("model", state.get("backbone", state))
|
||||
missing, unexpected = self.backbone.load_state_dict(weights, strict=False)
|
||||
logger.info(
|
||||
"已从 %s 加载骨架权重(缺失 %d 项,多余 %d 项,已训 %s 轮)",
|
||||
path,
|
||||
len(missing),
|
||||
len(unexpected),
|
||||
state.get("epochs", "?"),
|
||||
)
|
||||
if len(missing) > len(list(self.backbone.state_dict())) // 2:
|
||||
raise ValueError(
|
||||
f"{path} 与当前骨架结构对不上(缺失 {len(missing)} 项);"
|
||||
"确认 backbone 配置与自监督训练时一致"
|
||||
)
|
||||
|
||||
def forward(self, images: torch.Tensor) -> torch.Tensor:
|
||||
return self.head(self.backbone(images))
|
||||
Reference in New Issue
Block a user