diff --git a/configs/cluster/film.yaml b/configs/cluster/film.yaml new file mode 100644 index 0000000..67a46ea --- /dev/null +++ b/configs/cluster/film.yaml @@ -0,0 +1,109 @@ +# 未分类样本的无监督聚类。 +# +# 预处理决策（针对 data/film 的实际情况：14.5k 张已裁好的缺陷小图， +# 中位尺寸 152x113，宽高比 0.24~39.9，尺寸几乎每张都不同）： +# +# 前提：图片的尺寸和宽高比是上游 ROI 裁剪环节的产物，不是缺陷属性。 +# 所以预处理的目标是**把这个采集噪声抹掉**，只留下缺陷的纹理/形态。 +# +# 1. 直接 Resize 拉伸到 128x128，不用 letterbox —— letterbox 的填充比例 +# 恰好等价于编码了宽高比，灰边占多少就等于把采集噪声明着喂给网络； +# 拉伸则让所有图变成同一个画幅，宽高比不再是可用信号。 +# 2. use_shape_features 关掉 —— 同理，尺寸不进 embedding。 +# 3. 逐图标准化（PerImageStandardize）而不是固定的 ImageNet mean/std —— +# 整体亮度来自光照和相机增益，同样不是缺陷属性。用固定归一化时， +# 4 个簇里有 3 个花在"背景偏亮/偏暗"上；改成逐图 z-score 后， +# 簇变成"条纹状 / 团块状 / 弱纹理"，轮廓 0.256->0.322，稳定性 0.966->0.995。 +# 4. 不下采样到 224 —— 这批图本来就小，放大只会糊，128 已经够。 +# +# 代价：拉伸会让纹理各向异性形变（40:1 的长条尤其明显）。如果后续发现 +# 聚类被形变干扰，可选的替代是在原始分辨率上取固定大小的中心 patch +# （不缩放、不填充），配置里把 Resize 换成 CenterCrop 即可。 + +work_dir: runs/cluster +seed: 42 +device: cuda # method=pretrained 时可以 --set device=cpu + +num_clusters: 6 # 想分几类改这里（或 --num-clusters 6） + +# pretrained: ImageNet 预训练卷积骨架直接抽特征，不训练，CPU 也能跑，几分钟出结果。 +# simclr: 在本目录数据上自监督训练项目自己的卷积骨架，域适配更好，需要 Nvidia GPU。 +method: pretrained + +# method=pretrained 用的骨架（method=simclr 的骨架配在 simclr.backbone，两者字段完全不同， +# 各自独立以免 _base_ 合并时串味） +backbone: + name: torchvision + arch: resnet18 + pretrained: true + in_channels: 3 + +image_size: [128, 128] +use_shape_features: false # 尺寸/宽高比是采集产物，不作为聚类特征；true 仅供对照实验 + +# 取哪几个 stage 的特征拼起来（torchvision resnet 的 features 下标：4~7 = layer1~4）。 +# 实测（2000 张随机样本，k=4）： +# 只用最后一层 layer4 轮廓 0.191 换种子稳定性 0.410 +# layer2+3 轮廓 0.256 稳定性 0.989 <- 默认 +# layer2+3+4 轮廓 0.220 稳定性 0.627 +# 缺陷之间的差别主要在纹理（颗粒感、条纹方向、边缘锐度），中间层保留得比最后一层好； +# 最后一层的 ImageNet 语义（"这是猫还是车"）对这批图基本是噪声。 +feature_stages: [5, 6] + +# 特征降维。whiten 默认关闭：CNN 特征只有前若干主成分携带结构， +# 白化会把剩下几十个噪声方向抬到同权，簇立刻塌成"一个大簇 + 几个碎簇" +# （实测开白化时最大簇占 96%，关掉后降到 38%）。 +pca_dim: 32 +whiten: false +n_init: 10 + +# 只用一部分样本拟合质心以节省时间（0 = 用全部）；纯随机抽取，不按目录分层—— +# 目录名是人工整理的历史产物，不参与聚类。拟合完仍会把全部样本分配到最近的簇。 +fit_max_samples: 0 + +dataset: + name: unlabeled + root: data/film + gray: false + exclude: [images, labels] # 跳过已经整理成 YOLO 结构的部分 + transforms: + - {name: Resize, size: [128, 128]} # 拉伸，不保长宽比：见文件开头的说明 + - {name: ToTensor} + - {name: PerImageStandardize} # 逐图 z-score，抹掉整体亮度/对比度 + +loader: + batch_size: 64 + num_workers: 8 + +# method=simclr 时每个视图各自过一遍的随机增强（正样本对就是同一张图的两个视图） +view_transforms: + - {name: RandomResizedCrop, size: [128, 128], scale: [0.5, 1.0]} + - {name: RandomHorizontalFlip, p: 0.5} + - {name: RandomVerticalFlip, p: 0.5} + - {name: ColorJitter, brightness: 0.4, contrast: 0.4} + - {name: GaussianBlur, p: 0.3} + # 注意：主 transforms 里的 PerImageStandardize 会在这些增强之后执行， + # 所以 ColorJitter 的亮度扰动最终会被抵消掉，起作用的主要是裁剪/翻转/模糊。 + +simclr: + # 自监督用项目自己的纯卷积骨架：练出来的权重可以直接拿去初始化 DefectNet 的 backbone + backbone: + name: tiny_conv + in_channels: 3 + widths: [32, 64, 128, 256] + depths: [1, 2, 2, 2] + block: dw + feature_stages: null # tiny_conv 不是 nn.Sequential，取最终输出 + epochs: 100 + batch_size: 256 + temperature: 0.2 # 越小越关注难负样本 + amp: true + log_interval: 20 + checkpoint: null # 填已有的 simclr_backbone.pth 可跳过重复训练 + projection: {hidden_channels: 256, out_channels: 128} + optimizer: {name: adamw, lr: 0.001, weight_decay: 0.0001} + scheduler: {name: cosine} + +# 输出：每簇软链多少张图 / 预览拼图放多少张 +link_limit: 300 +preview_count: 25 diff --git a/datasets/unlabeled.py b/datasets/unlabeled.py new file mode 100644 index 0000000..a6a4f02 --- /dev/null +++ b/datasets/unlabeled.py @@ -0,0 +1,123 @@ +"""无标注图像数据集，用于无监督聚类。 + +递归扫描一个目录下的所有图片，不需要任何标注。 + +样本还会带出形状特征（宽、高、宽高比、面积的对数），但**默认不参与聚类**： +这些尺寸是上游 ROI 裁剪的产物，不是缺陷本身的属性，拿它做特征只会让聚类 +去区分采集方式。保留它只是为了写进 assignments.csv 供人工核对、以及做对照实验。 +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +import torch +from PIL import Image +from torch.utils.data import Dataset + +from datasets.registry import DATASETS +from datasets.transforms import build_transforms +from datasets.yolo import IMAGE_SUFFIXES + +SHAPE_FEATURE_DIM = 4 + + +def shape_features(width: int, height: int) -> torch.Tensor: + """[log 宽, log 高, log 宽高比, log 面积]，取对数是因为尺寸跨了两个数量级。""" + return torch.tensor( + [ + math.log(width), + math.log(height), + math.log(width / height), + math.log(width * height), + ], + dtype=torch.float32, + ) + + +@DATASETS.register("unlabeled") +class UnlabeledImageDataset(Dataset): + """Args: + root: 图片根目录，递归扫描。 + transforms: transform 配置列表；聚类默认用 Letterbox + ToTensor + Normalize。 + gray: 是否转单通道（RGBA / 调色板图会先统一到 RGB 或 L）。 + views: 每个样本返回几个增强视图。1 = 推理取特征；2 = SimCLR 对比学习。 + view_transforms: `views > 1` 时每个视图各自过一遍的随机增强。 + exclude: 需要跳过的目录名（默认跳过已划分好的 images/labels 与临时目录）。 + """ + + def __init__( + self, + root: str | Path, + transforms: Sequence[Mapping[str, Any]] = (), + gray: bool = False, + views: int = 1, + view_transforms: Sequence[Mapping[str, Any]] = (), + exclude: Sequence[str] = (), + ) -> None: + self.root = Path(root) + if not self.root.is_dir(): + raise FileNotFoundError(f"目录不存在: {self.root}") + self.gray = gray + self.views = int(views) + self.transforms = build_transforms(transforms) + self.view_transforms = build_transforms(view_transforms) if view_transforms else None + self.exclude = set(exclude) + + self.images = sorted( + path + for path in self.root.rglob("*") + if path.suffix.lower() in IMAGE_SUFFIXES + and not path.name.startswith(".") + and not self.exclude & set(path.relative_to(self.root).parts) + ) + if not self.images: + raise FileNotFoundError(f"{self.root} 下没有找到图片") + + def source_dir(self, path: Path) -> str: + """图片相对 root 的父目录，用来和聚类结果做交叉表（人工已分好的部分即是弱标签）。""" + relative = path.relative_to(self.root).parent + return str(relative) if str(relative) != "." else "(root)" + + def __len__(self) -> int: + return len(self.images) + + def __getitem__(self, index: int) -> dict: + path = self.images[index] + with Image.open(path) as raw: + image = raw.convert("L" if self.gray else "RGB") # RGBA / P 一并统一掉 + width, height = image.size + + if self.views == 1: + views = [self.transforms({"image": image.copy()})["image"]] + else: + views = [] + for _ in range(self.views): + sample = {"image": image.copy()} + if self.view_transforms is not None: + sample = self.view_transforms(sample) + views.append(self.transforms(sample)["image"]) + + return { + "views": views, + "shape": shape_features(width, height), + "meta": { + "path": str(path), + "index": index, + "source_dir": self.source_dir(path), + "size": (width, height), + }, + } + + +def unlabeled_collate(batch: list[dict]) -> dict: + """把每个视图各自堆成一个 batch：views[i] 形状 (B, C, H, W)。""" + num_views = len(batch[0]["views"]) + return { + "views": [torch.stack([item["views"][i] for item in batch]) for i in range(num_views)], + "shape": torch.stack([item["shape"] for item in batch]), + "meta": [item["meta"] for item in batch], + } diff --git a/engine/cluster.py b/engine/cluster.py new file mode 100644 index 0000000..ed90e1d --- /dev/null +++ b/engine/cluster.py @@ -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]) diff --git a/engine/simclr.py b/engine/simclr.py new file mode 100644 index 0000000..1e7075a --- /dev/null +++ b/engine/simclr.py @@ -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 diff --git a/models/heads/projection.py b/models/heads/projection.py new file mode 100644 index 0000000..0da3bb4 --- /dev/null +++ b/models/heads/projection.py @@ -0,0 +1,30 @@ +"""SimCLR 的投影头。 + +对比损失只作用在投影后的空间；聚类时**丢掉投影头、用骨架特征**—— +这是 SimCLR 的既有结论：投影头会丢掉对下游任务有用但对比任务用不上的信息。 +""" + +from __future__ import annotations + +import torch +from torch import nn + +from models.registry import HEADS + + +@HEADS.register("projection") +class ProjectionHead(nn.Module): + def __init__(self, in_channels: int, hidden_channels: int = 256, out_channels: int = 128) -> None: + super().__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.mlp = nn.Sequential( + nn.Linear(in_channels, hidden_channels), + nn.BatchNorm1d(hidden_channels), + nn.ReLU(inplace=True), + nn.Linear(hidden_channels, out_channels), + ) + + def forward(self, feat: torch.Tensor) -> torch.Tensor: + if feat.dim() == 4: + feat = self.pool(feat).flatten(1) + return self.mlp(feat) diff --git a/scripts/cluster_unlabeled.py b/scripts/cluster_unlabeled.py new file mode 100644 index 0000000..54b4b2c --- /dev/null +++ b/scripts/cluster_unlabeled.py @@ -0,0 +1,296 @@ +#!/usr/bin/env python +"""把未分类的样本做无监督聚类，产出可人工核对的分组结果。 + + # 默认 4 类，ImageNet 预训练特征（不训练，CPU 可跑） + uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml + + # 改簇数 / 换成在本地数据上自监督训练的特征（需要 Nvidia GPU） + uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \ + --set num_clusters=6 method=simclr + +产出（work_dir 下）： + assignments.csv 每张图的簇号、到质心距离、原目录 + summary.json 簇大小、轮廓系数、簇 × 原目录交叉表 + clusters/cluster_00/... 按簇建的软链目录，直接用看图软件翻 + cluster_00_preview.png 每簇代表样本（离质心最近）的拼图 +""" + +from __future__ import annotations + +import argparse +import csv +import json +import shutil +import sys +from collections import defaultdict +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import torch # noqa: E402 + +from datasets.build import build_dataset # noqa: E402 +from datasets.unlabeled import unlabeled_collate # noqa: E402 +from engine.cluster import ( # noqa: E402 + StageFeatureExtractor, + cluster_features, + cluster_stability, + extract_features, + random_fit_index, +) +from models.registry import BACKBONES # noqa: E402 +from utils import load_config, seed_everything, setup_logging # noqa: E402 +from utils.device import NoNvidiaGpuError, require_cuda, resolve_device # noqa: E402 + +logger = None + + +def build_encoder(cfg, device, work_dir): + """按 `method` 拿一个特征提取器。 + + pretrained: ImageNet 预训练卷积骨架，纯推理，CPU 也能跑。用顶层的 `backbone`。 + simclr: 在本目录数据上自监督训练一个项目自己的卷积骨架——属于训练，先过 GPU 守卫。 + 用 `simclr.backbone`，与 pretrained 的骨架配置各自独立：两者的字段完全不同 + （arch/pretrained vs widths/depths），共用一个块会在 `_base_` 合并时串味。 + """ + method = str(cfg.get("method", "pretrained")) + + if method == "pretrained": + backbone = BACKBONES.build(dict(cfg.get("backbone", {"name": "torchvision"}))) + return StageFeatureExtractor(backbone, cfg.get("feature_stages")).to(device).eval() + + if method != "simclr": + raise ValueError(f"未知 method: {method}，可选 pretrained / simclr") + + require_cuda() # 自监督预训练是训练，禁止在无 Nvidia GPU 的设备上跑 + from torch.utils.data import DataLoader + + from engine.simclr import SimCLRTrainer + + simclr_cfg = dict(cfg.get("simclr", {})) + backbone_cfg = dict(simclr_cfg.pop("backbone", {"name": "tiny_conv"})) + # 自研骨架不是 nn.Sequential，默认取最终输出而不是中间 stage + stages = simclr_cfg.pop("feature_stages", None) + backbone = BACKBONES.build(backbone_cfg).to(device) + + checkpoint = simclr_cfg.pop("checkpoint", None) + if checkpoint and Path(checkpoint).exists(): + state = torch.load(checkpoint, map_location=device, weights_only=False) + backbone.load_state_dict(state["model"]) + logger.info("复用已有的自监督骨架: %s", checkpoint) + return StageFeatureExtractor(backbone, stages).to(device).eval() + + view_dataset = build_dataset( + {**dict(cfg["dataset"]), "views": 2, "view_transforms": cfg.get("view_transforms", [])} + ) + loader = DataLoader( + view_dataset, + batch_size=int(simclr_cfg.pop("batch_size", 128)), + shuffle=True, + drop_last=True, + num_workers=int(cfg.get("loader", {}).get("num_workers", 4)), + collate_fn=unlabeled_collate, + ) + SimCLRTrainer(backbone, loader, simclr_cfg, device, work_dir).fit() + return StageFeatureExtractor(backbone, stages).to(device).eval() + + +def write_outputs(cfg, work_dir, metas, result, num_clusters, stability): + """CSV + JSON + 软链目录 + 预览拼图。人工只需要翻软链目录就能确认分组是否合理。""" + labels = result.labels.tolist() + distances = result.distances.tolist() + + with (work_dir / "assignments.csv").open("w", newline="", encoding="utf-8") as handle: + writer = csv.writer(handle) + writer.writerow(["path", "cluster", "distance_to_centroid", "source_dir", "width", "height"]) + for meta, label, distance in zip(metas, labels, distances, strict=True): + width, height = meta["size"] + writer.writerow([meta["path"], label, f"{distance:.6f}", meta["source_dir"], width, height]) + + summary = { + "num_clusters": num_clusters, + "num_images": len(metas), + "method": cfg.get("method", "pretrained"), + "image_size": cfg.get("image_size"), + "cluster_sizes": result.sizes, + "silhouette": result.silhouette, + "stability": stability, + "inertia": result.inertia, + "mean_distance_per_cluster": [ + float(result.distances[result.labels == index].mean()) if size else None + for index, size in enumerate(result.sizes) + ], + } + (work_dir / "summary.json").write_text( + json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8" + ) + + link_limit = int(cfg.get("link_limit", 300)) + preview_count = int(cfg.get("preview_count", 25)) + clusters_dir = work_dir / "clusters" + if clusters_dir.exists(): + shutil.rmtree(clusters_dir) + + order = sorted(range(len(metas)), key=lambda i: distances[i]) # 离质心近的排前面 + per_cluster: dict[int, list[int]] = defaultdict(list) + for index in order: + per_cluster[labels[index]].append(index) + + for cluster, indices in sorted(per_cluster.items()): + target = clusters_dir / f"cluster_{cluster:02d}" + target.mkdir(parents=True, exist_ok=True) + # 近的取大头，末尾补几张最远的——离群样本最能暴露这个簇是不是混了东西 + chosen = indices[: max(link_limit - 20, 1)] + indices[-20:] if len(indices) > link_limit else indices + for rank, index in enumerate(chosen): + source = Path(metas[index]["path"]).resolve() + link = target / f"{rank:04d}_{source.stem}{source.suffix}" + if not link.exists(): + link.symlink_to(source) + save_preview(work_dir / f"cluster_{cluster:02d}_preview.png", metas, indices[:preview_count]) + + return summary + + +def save_preview(path: Path, metas: list[dict], indices: list[int], cell: int = 96, columns: int = 5): + """把代表样本拼成一张图，方便一眼看出这个簇是什么缺陷。""" + if not indices: + return + from PIL import Image + + rows = (len(indices) + columns - 1) // columns + canvas = Image.new("RGB", (columns * cell, rows * cell), (32, 32, 32)) + for position, index in enumerate(indices): + with Image.open(metas[index]["path"]) as raw: + image = raw.convert("RGB") + scale = min(cell / image.width, cell / image.height) + image = image.resize((max(1, round(image.width * scale)), max(1, round(image.height * scale)))) + x = (position % columns) * cell + (cell - image.width) // 2 + y = (position // columns) * cell + (cell - image.height) // 2 + canvas.paste(image, (x, y)) + canvas.save(path) + + +def sweep_clusters(embeddings, spec: str, cluster_kwargs: dict) -> None: + """扫一遍簇数，同时看轮廓系数和换种子的稳定性。 + + 无标签时没有"正确答案"，只能看：簇是否紧致（轮廓），以及换初始化后划分是否还一样 + （稳定性）。稳定性低说明那个 k 是在硬切一团连续的数据。 + """ + start, end = (int(v) for v in spec.split("-")) + logger.info("%3s %9s %10s %12s 簇大小", "k", "轮廓", "稳定性", "最大簇占比") + for k in range(start, end + 1): + result = cluster_features(embeddings, k, seed=42, **cluster_kwargs) + stability = cluster_stability(embeddings, k, seeds=(0, 7), **cluster_kwargs) + logger.info( + "%3d %9.3f %10.3f %12.2f %s", + k, + result.silhouette, + stability, + max(result.sizes) / sum(result.sizes), + result.sizes, + ) + + +def main() -> int: + global logger + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--config", default="configs/cluster/film.yaml") + parser.add_argument("--root", default=None, help="覆盖配置里的图片根目录") + parser.add_argument("--num-clusters", type=int, default=None, help="覆盖配置里的簇数") + parser.add_argument("--work-dir", default=None) + parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE") + parser.add_argument("--save-features", action="store_true", help="额外保存 embedding 供复用") + parser.add_argument( + "--sweep-k", default=None, metavar="A-B", help="只扫簇数（如 2-8）打印指标表，不产出分组" + ) + args = parser.parse_args() + + cfg = load_config(args.config, args.set) + if args.root: + cfg["dataset"]["root"] = args.root + if args.num_clusters: + cfg["num_clusters"] = args.num_clusters + + work_dir = Path(args.work_dir or cfg.get("work_dir", "runs/cluster")) + work_dir.mkdir(parents=True, exist_ok=True) + logger = setup_logging(work_dir / "cluster.log") + seed_everything(int(cfg.get("seed", 42))) + + num_clusters = int(cfg.get("num_clusters", 4)) + try: + device = resolve_device(str(cfg.get("device", "cuda"))) + except NoNvidiaGpuError as exc: + logger.error("%s", exc) + return 1 + + dataset = build_dataset(dict(cfg["dataset"])) + logger.info( + "扫描到 %d 张图，来自 %d 个目录；簇数 %d，method=%s，device=%s", + len(dataset), + len({dataset.source_dir(path) for path in dataset.images}), + num_clusters, + cfg.get("method", "pretrained"), + device, + ) + + try: + encoder = build_encoder(cfg, device, work_dir) + except NoNvidiaGpuError as exc: + logger.error("%s", exc) + return 1 + + from torch.utils.data import DataLoader + + loader = DataLoader( + dataset, + batch_size=int(cfg.get("loader", {}).get("batch_size", 64)), + shuffle=False, + num_workers=int(cfg.get("loader", {}).get("num_workers", 4)), + collate_fn=unlabeled_collate, + ) + embeddings, metas = extract_features( + encoder, loader, device, use_shape_features=bool(cfg.get("use_shape_features", False)) + ) + logger.info("embedding: %s", tuple(embeddings.shape)) + + cluster_kwargs = { + "pca_dim": int(cfg.get("pca_dim", 32)), + "whiten": bool(cfg.get("whiten", False)), + "n_init": int(cfg.get("n_init", 10)), + } + fit_index = random_fit_index( + len(metas), int(cfg.get("fit_max_samples", 0)), seed=int(cfg.get("seed", 42)) + ) + if fit_index is not None: + logger.info("随机抽 %d 个样本拟合质心（拟合后仍分配全部样本）", len(fit_index)) + + if args.sweep_k: + sweep_clusters(embeddings, args.sweep_k, cluster_kwargs) + return 0 + + result = cluster_features( + embeddings, num_clusters, seed=int(cfg.get("seed", 42)), fit_index=fit_index, **cluster_kwargs + ) + stability = cluster_stability(embeddings, num_clusters, seeds=(0, 7), **cluster_kwargs) + logger.info( + "簇大小: %s | 轮廓系数: %.4f | 换种子稳定性: %.4f", + result.sizes, + result.silhouette, + stability, + ) + + if args.save_features: + torch.save({"embeddings": embeddings, "metas": metas}, work_dir / "features.pt") + + write_outputs(cfg, work_dir, metas, result, num_clusters, stability) + logger.info("结果已写入 %s（clusters/ 下是按簇建的软链目录，直接翻图即可核对）", work_dir) + for index, size in enumerate(result.sizes): + mean_distance = float(result.distances[result.labels == index].mean()) if size else float("nan") + logger.info(" cluster_%02d %6d 张 平均距质心 %.3f", index, size, mean_distance) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())