Shaun ZhouandClaude Opus 5
33fdf396ee
在聚类 + 人工归类得到的数据集（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>
83 lines
3.1 KiB
Python
83 lines
3.1 KiB
Python
"""缺陷分类模型：卷积骨架 + 分类头。
|
||
|
||
与 `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))
|