feat: 未分类样本的无监督聚类（多 stage 特征 + KMeans，可选 SimCLR）
把 data/film 下所有图片拉平做聚类，不使用任何目录名，簇数由 num_clusters 决定。 预处理针对该批数据的实际情况（14.5k 张已裁好的缺陷小图，中位 152x113， 宽高比 0.24~39.9）：尺寸与亮度都是 ROI 裁剪和光照的产物、不是缺陷属性， 因此拉伸到 128x128（不用 letterbox，其填充比例等价于编码了宽高比）并逐图 z-score。 - cluster: torch 实现的 kmeans++/Lloyd（含空簇处理与 n_init）、PCA、轮廓系数、 NMI 稳定性、多 stage 特征提取；不引 scikit-learn - 白化默认关闭：CNN 特征谱极陡（λ1/λ32≈165x），白化会把噪声方向抬到与信号同权， 实测最大簇一度吃掉 96% 的样本 - simclr: NT-Xent 自监督预训练，用项目自己的卷积骨架，权重可直接初始化 DefectNet； 属于训练，走 require_cuda() 守卫 - 产出 assignments.csv、summary.json、按簇的软链目录与代表样本拼图，供人工核对 全量 14,579 张：轮廓系数 0.326、换种子稳定性 0.964，CPU 约 100 秒。 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
4b1147ed8f
commit
f433e5cf59
@@ -0,0 +1,293 @@
|
||||
"""无监督聚类：特征提取 + KMeans。
|
||||
|
||||
用 torch 自己实现 KMeans（kmeans++ 初始化 + Lloyd 迭代），不引 scikit-learn：
|
||||
一是少一个依赖，二是能直接吃 GPU 上的特征张量。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from utils.logging import get_logger
|
||||
|
||||
logger = get_logger("engine.cluster")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClusterResult:
|
||||
labels: torch.Tensor # (N,) 每个样本的簇下标
|
||||
centroids: torch.Tensor # (K, D)
|
||||
distances: torch.Tensor # (N,) 到所属质心的距离
|
||||
inertia: float # 簇内平方距离和，越小越紧
|
||||
silhouette: float = float("nan") # 轮廓系数，[-1, 1]，越大越好
|
||||
sizes: list[int] = field(default_factory=list)
|
||||
|
||||
|
||||
def l2_normalize(x: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
|
||||
return x / x.norm(dim=1, keepdim=True).clamp(min=eps)
|
||||
|
||||
|
||||
def pca_reduce(x: torch.Tensor, dim: int, whiten: bool = False, eps: float = 1e-5) -> torch.Tensor:
|
||||
"""PCA 降维，`whiten` 决定是否把各主成分方差拉平。
|
||||
|
||||
**白化默认关闭**：CNN 特征里只有前若干个主成分携带结构，其余接近噪声；
|
||||
白化会把这些噪声方向抬到和信号同权，簇随即塌成"一个大簇 + 几个碎簇"。
|
||||
实测 3 个信号维 + 200 个噪声维的合成数据上，dim=3 轮廓 0.93，
|
||||
dim=64 + 白化掉到 0.10 且最大簇占 98%。需要白化时请同时把 dim 调小。
|
||||
"""
|
||||
dim = min(dim, x.shape[0], x.shape[1])
|
||||
centered = x - x.mean(dim=0, keepdim=True)
|
||||
_, singular, components = torch.pca_lowrank(centered, q=dim, center=False)
|
||||
projected = centered @ components[:, :dim]
|
||||
if not whiten:
|
||||
return projected
|
||||
return projected / (singular[:dim] / (x.shape[0] ** 0.5) + eps)
|
||||
|
||||
|
||||
def kmeans_plusplus(x: torch.Tensor, k: int, generator: torch.Generator) -> torch.Tensor:
|
||||
"""kmeans++ 初始化：按到已选质心距离的平方为概率挑下一个，避免初值退化。"""
|
||||
n = x.shape[0]
|
||||
centroids = [x[torch.randint(n, (1,), generator=generator, device=x.device).item()]]
|
||||
closest = torch.cdist(x, centroids[0][None]).squeeze(1) ** 2
|
||||
for _ in range(1, k):
|
||||
probs = closest / closest.sum().clamp(min=1e-12)
|
||||
index = torch.multinomial(probs, 1, generator=generator).item()
|
||||
centroids.append(x[index])
|
||||
closest = torch.minimum(closest, torch.cdist(x, centroids[-1][None]).squeeze(1) ** 2)
|
||||
return torch.stack(centroids)
|
||||
|
||||
|
||||
def kmeans(
|
||||
x: torch.Tensor,
|
||||
num_clusters: int,
|
||||
iters: int = 100,
|
||||
n_init: int = 10,
|
||||
tol: float = 1e-5,
|
||||
seed: int = 42,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, float]:
|
||||
"""跑 `n_init` 次取 inertia 最小的一次，返回 (labels, centroids, inertia)。"""
|
||||
if num_clusters < 2:
|
||||
raise ValueError(f"num_clusters 至少为 2，收到 {num_clusters}")
|
||||
if x.shape[0] < num_clusters:
|
||||
raise ValueError(f"样本数 {x.shape[0]} 少于簇数 {num_clusters}")
|
||||
|
||||
best: tuple[torch.Tensor, torch.Tensor, float] | None = None
|
||||
for run in range(n_init):
|
||||
generator = torch.Generator(device=x.device).manual_seed(seed + run)
|
||||
centroids = kmeans_plusplus(x, num_clusters, generator)
|
||||
previous = float("inf")
|
||||
|
||||
for _ in range(iters):
|
||||
distances = torch.cdist(x, centroids)
|
||||
labels = distances.argmin(dim=1)
|
||||
inertia = float((distances.gather(1, labels[:, None]) ** 2).sum())
|
||||
|
||||
for index in range(num_clusters):
|
||||
members = x[labels == index]
|
||||
if len(members):
|
||||
centroids[index] = members.mean(dim=0)
|
||||
else:
|
||||
# 空簇：把离自己质心最远的点抢过来，避免簇数缩水
|
||||
centroids[index] = x[distances.gather(1, labels[:, None]).squeeze(1).argmax()]
|
||||
if abs(previous - inertia) <= tol * max(previous, 1.0):
|
||||
break
|
||||
previous = inertia
|
||||
|
||||
if best is None or inertia < best[2]:
|
||||
best = (labels, centroids.clone(), inertia)
|
||||
|
||||
assert best is not None
|
||||
return best
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def silhouette_score(
|
||||
x: torch.Tensor, labels: torch.Tensor, max_samples: int = 2000, seed: int = 42
|
||||
) -> float:
|
||||
"""轮廓系数。是 O(N²) 的，所以超过 `max_samples` 就抽样估计。"""
|
||||
n = x.shape[0]
|
||||
if labels.unique().numel() < 2:
|
||||
return float("nan")
|
||||
if n > max_samples:
|
||||
generator = torch.Generator(device="cpu").manual_seed(seed)
|
||||
index = torch.randperm(n, generator=generator)[:max_samples].to(x.device)
|
||||
x, labels = x[index], labels[index]
|
||||
|
||||
distances = torch.cdist(x, x)
|
||||
scores = []
|
||||
for index in range(x.shape[0]):
|
||||
same = labels == labels[index]
|
||||
same[index] = False
|
||||
if not same.any():
|
||||
continue # 单点簇按惯例记 0
|
||||
a = distances[index][same].mean()
|
||||
b = min(
|
||||
distances[index][labels == other].mean()
|
||||
for other in labels.unique().tolist()
|
||||
if other != labels[index].item()
|
||||
)
|
||||
scores.append(((b - a) / torch.maximum(a, b)).item())
|
||||
return float(sum(scores) / len(scores)) if scores else float("nan")
|
||||
|
||||
|
||||
class StageFeatureExtractor(nn.Module):
|
||||
"""从卷积骨架的多个 stage 取特征，各自全局池化 + L2 归一化后拼接。
|
||||
|
||||
只用最后一层的语义特征对薄膜缺陷不够用——缺陷之间的差异很大程度是**纹理**层面的
|
||||
（颗粒感、条纹方向、边缘锐度），这些信息在中间层比在最后一层保留得好。
|
||||
"""
|
||||
|
||||
def __init__(self, backbone: nn.Module, stages: Sequence[int] | None = None) -> None:
|
||||
super().__init__()
|
||||
self.backbone = backbone
|
||||
layers = getattr(backbone, "features", backbone)
|
||||
self.layers = list(layers) if isinstance(layers, nn.Sequential) else None
|
||||
self.stages = sorted(stages) if stages else None
|
||||
if self.stages and self.layers is None:
|
||||
raise ValueError("该骨架不是 nn.Sequential，无法按 stage 取中间特征")
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if not self.stages:
|
||||
feat = self.backbone(x)
|
||||
return feat.mean(dim=(2, 3)) if feat.dim() == 4 else feat
|
||||
|
||||
outputs = []
|
||||
for index, layer in enumerate(self.layers):
|
||||
x = layer(x)
|
||||
if index in self.stages:
|
||||
outputs.append(l2_normalize(x.mean(dim=(2, 3))))
|
||||
if not outputs:
|
||||
raise ValueError(f"stages {self.stages} 超出骨架层数 {len(self.layers)}")
|
||||
return torch.cat(outputs, dim=1)
|
||||
|
||||
|
||||
def normalized_mutual_info(labels_a: Sequence, labels_b: Sequence) -> float:
|
||||
"""两组划分的归一化互信息，[0, 1]。这里用来比较**两次聚类结果**的一致性。"""
|
||||
from collections import Counter
|
||||
|
||||
n = len(labels_a)
|
||||
if n == 0:
|
||||
return float("nan")
|
||||
count_a, count_b = Counter(labels_a), Counter(labels_b)
|
||||
joint = Counter(zip(labels_a, labels_b, strict=True))
|
||||
|
||||
def entropy(counter: Counter) -> float:
|
||||
return -sum((c / n) * math.log(c / n) for c in counter.values() if c)
|
||||
|
||||
mutual = sum(
|
||||
(c / n) * math.log((c / n) / ((count_a[a] / n) * (count_b[b] / n)))
|
||||
for (a, b), c in joint.items()
|
||||
if c
|
||||
)
|
||||
denominator = (entropy(count_a) * entropy(count_b)) ** 0.5
|
||||
return float(mutual / denominator) if denominator > 0 else float("nan")
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def extract_features(
|
||||
encoder: nn.Module,
|
||||
dataloader: DataLoader,
|
||||
device: torch.device,
|
||||
use_shape_features: bool = False,
|
||||
log_interval: int = 20,
|
||||
) -> tuple[torch.Tensor, list[dict]]:
|
||||
"""跑一遍数据集拿 embedding。
|
||||
|
||||
`encoder` 输出特征图，这里做全局平均池化。
|
||||
|
||||
`use_shape_features` 默认关闭：图片的宽高比/尺寸是 ROI 裁剪环节的产物，
|
||||
不是缺陷属性，拼进 embedding 只会让聚类去区分"采集方式"而不是缺陷类型。
|
||||
留着这个开关只为做对照实验。
|
||||
"""
|
||||
encoder.eval()
|
||||
features: list[torch.Tensor] = []
|
||||
shapes: list[torch.Tensor] = []
|
||||
metas: list[dict] = []
|
||||
|
||||
for step, batch in enumerate(dataloader):
|
||||
images = batch["views"][0].to(device, non_blocking=True)
|
||||
feat = encoder(images)
|
||||
if feat.dim() == 4:
|
||||
feat = feat.mean(dim=(2, 3))
|
||||
features.append(feat.float().cpu())
|
||||
shapes.append(batch["shape"])
|
||||
metas.extend(batch["meta"])
|
||||
if step % log_interval == 0:
|
||||
logger.info("提取特征 [%d/%d]", step, len(dataloader))
|
||||
|
||||
embeddings = torch.cat(features)
|
||||
if use_shape_features:
|
||||
shape = torch.cat(shapes)
|
||||
shape = (shape - shape.mean(dim=0)) / shape.std(dim=0).clamp(min=1e-6)
|
||||
# 图像特征先归一化，避免两部分量纲差太多导致形状特征被淹没
|
||||
embeddings = torch.cat([l2_normalize(embeddings), shape], dim=1)
|
||||
return embeddings, metas
|
||||
|
||||
|
||||
def cluster_features(
|
||||
embeddings: torch.Tensor,
|
||||
num_clusters: int,
|
||||
pca_dim: int = 32,
|
||||
whiten: bool = False,
|
||||
normalize: bool = True,
|
||||
n_init: int = 10,
|
||||
iters: int = 100,
|
||||
seed: int = 42,
|
||||
fit_index: torch.Tensor | None = None,
|
||||
) -> ClusterResult:
|
||||
"""L2 归一化 -> PCA 降维 -> KMeans -> 指标。
|
||||
|
||||
`fit_index` 给定时只用这部分样本拟合质心（数据极不均衡时用来做类别平衡采样），
|
||||
再把全部样本按最近质心分配。
|
||||
"""
|
||||
x = l2_normalize(embeddings) if normalize else embeddings
|
||||
if pca_dim:
|
||||
x = pca_reduce(x, pca_dim, whiten=whiten)
|
||||
|
||||
fit_x = x if fit_index is None else x[fit_index]
|
||||
_, centroids, _ = kmeans(fit_x, num_clusters, iters=iters, n_init=n_init, seed=seed)
|
||||
|
||||
distances = torch.cdist(x, centroids)
|
||||
labels = distances.argmin(dim=1)
|
||||
nearest = distances.gather(1, labels[:, None]).squeeze(1)
|
||||
|
||||
return ClusterResult(
|
||||
labels=labels,
|
||||
centroids=centroids,
|
||||
distances=nearest,
|
||||
inertia=float((nearest**2).sum()),
|
||||
silhouette=silhouette_score(x, labels, seed=seed),
|
||||
sizes=torch.bincount(labels, minlength=num_clusters).tolist(),
|
||||
)
|
||||
|
||||
|
||||
def random_fit_index(num_samples: int, max_samples: int, seed: int = 42) -> torch.Tensor | None:
|
||||
"""随机抽一部分样本用于拟合质心，纯粹为了省时间（拟合完仍会分配全部样本）。
|
||||
|
||||
不做任何按来源目录的分层——目录名是人工整理的历史产物，不应该影响聚类。
|
||||
"""
|
||||
if not max_samples or num_samples <= max_samples:
|
||||
return None
|
||||
generator = torch.Generator().manual_seed(seed)
|
||||
return torch.randperm(num_samples, generator=generator)[:max_samples].sort().values
|
||||
|
||||
|
||||
def cluster_stability(
|
||||
embeddings: torch.Tensor, num_clusters: int, seeds: Sequence[int] = (0, 1), **kwargs
|
||||
) -> float:
|
||||
"""换随机种子重跑，看两次划分的 NMI。
|
||||
|
||||
无标签场景下判断簇数是否合理的实用指标：真有结构时不同初始化会收敛到同一套划分，
|
||||
强行切分一团连续数据时结果会随种子乱跳。
|
||||
"""
|
||||
runs = [
|
||||
cluster_features(embeddings, num_clusters, seed=seed, **kwargs).labels.tolist()
|
||||
for seed in seeds
|
||||
]
|
||||
return normalized_mutual_info(runs[0], runs[1])
|
||||
@@ -0,0 +1,125 @@
|
||||
"""SimCLR 自监督预训练：在无标注的薄膜图上学一套特征，供聚类使用。
|
||||
|
||||
骨架仍然是项目里的纯卷积 backbone（`tiny_conv` 等），预训练完可以直接拿去
|
||||
初始化检测模型。属于"训练"，因此调用方必须先过 `utils.device.require_cuda()`。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from engine.metrics import AverageMeter
|
||||
from engine.optim import build_optimizer, build_scheduler
|
||||
from models.heads.projection import ProjectionHead
|
||||
from utils.logging import get_logger
|
||||
|
||||
logger = get_logger("engine.simclr")
|
||||
|
||||
|
||||
def nt_xent_loss(z1: torch.Tensor, z2: torch.Tensor, temperature: float = 0.2) -> torch.Tensor:
|
||||
"""NT-Xent：同一张图的两个视图互为正样本，batch 内其余 2N-2 个都是负样本。"""
|
||||
batch_size = z1.shape[0]
|
||||
z = F.normalize(torch.cat([z1, z2]), dim=1)
|
||||
similarity = z @ z.t() / temperature
|
||||
similarity.fill_diagonal_(float("-inf")) # 排除自己和自己
|
||||
|
||||
targets = torch.cat(
|
||||
[
|
||||
torch.arange(batch_size, 2 * batch_size, device=z.device),
|
||||
torch.arange(0, batch_size, device=z.device),
|
||||
]
|
||||
)
|
||||
return F.cross_entropy(similarity, targets)
|
||||
|
||||
|
||||
class SimCLRTrainer:
|
||||
"""Args:
|
||||
backbone: 输出特征图的卷积骨架（需带 `out_channels`）。
|
||||
dataloader: 每个样本产出 2 个视图的 DataLoader（`UnlabeledImageDataset(views=2)`）。
|
||||
cfg: `epochs / temperature / optimizer / scheduler / amp / log_interval` 等。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
backbone: nn.Module,
|
||||
dataloader: DataLoader,
|
||||
cfg: Mapping[str, Any],
|
||||
device: torch.device,
|
||||
work_dir: str | Path,
|
||||
) -> None:
|
||||
self.backbone = backbone
|
||||
self.dataloader = dataloader
|
||||
self.cfg = dict(cfg)
|
||||
self.device = device
|
||||
self.work_dir = Path(work_dir)
|
||||
self.work_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.epochs = int(self.cfg.get("epochs", 50))
|
||||
self.temperature = float(self.cfg.get("temperature", 0.2))
|
||||
self.log_interval = int(self.cfg.get("log_interval", 20))
|
||||
|
||||
self.projector = ProjectionHead(
|
||||
backbone.out_channels, **dict(self.cfg.get("projection", {}))
|
||||
).to(device)
|
||||
self.model = nn.ModuleDict({"backbone": backbone, "projector": self.projector})
|
||||
self.optimizer = build_optimizer(
|
||||
self.model, self.cfg.get("optimizer", {"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4})
|
||||
)
|
||||
self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
|
||||
self.amp = bool(self.cfg.get("amp", True)) and device.type == "cuda"
|
||||
self.scaler = torch.amp.GradScaler("cuda", enabled=self.amp)
|
||||
|
||||
def train_one_epoch(self, epoch: int) -> float:
|
||||
self.model.train()
|
||||
meter = AverageMeter()
|
||||
start = time.time()
|
||||
|
||||
for step, batch in enumerate(self.dataloader):
|
||||
view1, view2 = (view.to(self.device, non_blocking=True) for view in batch["views"][:2])
|
||||
|
||||
with torch.amp.autocast("cuda", enabled=self.amp):
|
||||
z1 = self.projector(self.backbone(view1))
|
||||
z2 = self.projector(self.backbone(view2))
|
||||
loss = nt_xent_loss(z1, z2, self.temperature)
|
||||
|
||||
self.optimizer.zero_grad(set_to_none=True)
|
||||
self.scaler.scale(loss).backward()
|
||||
self.scaler.step(self.optimizer)
|
||||
self.scaler.update()
|
||||
|
||||
meter.update(float(loss.detach()), view1.shape[0])
|
||||
if step % self.log_interval == 0:
|
||||
logger.info(
|
||||
"simclr epoch %d [%d/%d] loss=%.4f lr=%.2e",
|
||||
epoch,
|
||||
step,
|
||||
len(self.dataloader),
|
||||
meter.avg,
|
||||
self.optimizer.param_groups[0]["lr"],
|
||||
)
|
||||
|
||||
logger.info("simclr epoch %d 完成 loss=%.4f 耗时 %.1fs", epoch, meter.avg, time.time() - start)
|
||||
return meter.avg
|
||||
|
||||
def fit(self) -> Path:
|
||||
"""训练并保存骨架权重（不含投影头），返回权重路径。"""
|
||||
if len(self.dataloader) < 2:
|
||||
raise ValueError("SimCLR 需要多个 batch，请调小 batch_size 或增加样本")
|
||||
|
||||
for epoch in range(self.epochs):
|
||||
self.train_one_epoch(epoch)
|
||||
if self.scheduler is not None:
|
||||
self.scheduler.step()
|
||||
|
||||
path = self.work_dir / "simclr_backbone.pth"
|
||||
torch.save({"model": self.backbone.state_dict(), "epochs": self.epochs}, path)
|
||||
logger.info("自监督骨架已保存: %s", path)
|
||||
return path
|
||||
Reference in New Issue
Block a user