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,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