Files
ModelTrainEval/scripts/cluster_unlabeled.py
T
Shaun ZhouandClaude Opus 5 0f5bb4a5c0 feat(cluster): 加 --train-only,自监督训练与聚类可分两步跑
engine/simclr.py 与 engine/cluster.py 本来就互不引用,唯一接口是
simclr_backbone.pth;耦合只在 CLI——method=simclr 会训练完直接聚类。
此前"只聚类"可以用 simclr.checkpoint 做到,"只训练"没有办法。

- --train-only:训练存完权重即退出,不抽特征也不聚类,并打印下一步的聚类命令
- 与 method=pretrained / simclr.checkpoint / --sweep-k 组合时立刻报错说明原因,
  而不是跑完才发现没有产出
- 典型用法:GPU 机器上 --train-only 跑几十小时,之后拿权重反复调 k,每次约 1 分钟

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-07 11:38:59 +08:00

510 lines
20 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python
"""把未分类的样本做无监督聚类,产出可人工核对的分组结果。
# 默认 4 类,ImageNet 预训练特征(不训练,CPU 可跑)
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml
# 改簇数 / 换成在本地数据上自监督训练的特征(Mac 上自动走 MPS)
# 断点续训:--set method=simclr simclr.resume=runs/simclr/simclr_last.pth
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,
centroid_distances,
cluster_features,
cluster_stability,
extract_features,
project_2d,
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 ( # noqa: E402
DeviceNotAvailableError,
describe_device,
resolve_device,
)
logger = None
def build_encoder(cfg, device, work_dir):
"""按 `method` 拿一个特征提取器。
pretrained: ImageNet 预训练卷积骨架,纯推理,CPU 也能跑。用顶层的 `backbone`。
simclr: 在本目录数据上自监督训练一个项目自己的卷积骨架。
用 `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")
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)
resume = simclr_cfg.pop("resume", None)
checkpoint = simclr_cfg.pop("checkpoint", None)
if resume and checkpoint:
raise ValueError(
"simclr.resume 与 simclr.checkpoint 只能二选一:"
"resume=接着训练,checkpoint=跳过训练直接聚类。"
)
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,
)
full_config = cfg.to_dict() if hasattr(cfg, "to_dict") else dict(cfg)
SimCLRTrainer(
backbone, loader, simclr_cfg, device, work_dir, full_config=full_config, resume=resume
).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])
method = str(cfg.get("method", "pretrained"))
backbone_cfg = (
dict(cfg.get("simclr", {}).get("backbone", {}))
if method == "simclr"
else dict(cfg.get("backbone", {}))
)
summary = {
"num_clusters": num_clusters,
"num_images": len(metas),
"method": method,
# 记录特征来源,供 export_clusters.py 自动命名以及日后追溯
"feature_source": {
"method": method,
"backbone": backbone_cfg,
"feature_stages": (
cfg.get("simclr", {}).get("feature_stages")
if method == "simclr"
else cfg.get("feature_stages")
),
},
"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)
],
"centroid_distances": centroid_distances(result.centroids).tolist(),
}
(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])
if cfg.get("scatter", {}).get("enabled", True) and result.points is not None:
scatter_path = work_dir / "cluster_scatter.png"
save_scatter(scatter_path, result, cfg)
logger.info("簇分布图已保存: %s", scatter_path)
return summary
#: 按优先级尝试的中文字体(macOS / Linux 常见)。都找不到就退回英文标签,避免画出豆腐块。
CJK_FONTS = (
"PingFang SC",
"Hiragino Sans GB",
"Heiti SC",
"Noto Sans CJK SC",
"Source Han Sans SC",
"WenQuanYi Zen Hei",
"Microsoft YaHei",
"SimHei",
"Arial Unicode MS",
)
def setup_cjk_font() -> bool:
"""给 matplotlib 配一个能显示中文的字体,返回是否配上了。"""
import matplotlib
from matplotlib import font_manager
available = {f.name for f in font_manager.fontManager.ttflist}
for name in CJK_FONTS:
if name in available:
matplotlib.rcParams["font.sans-serif"] = [name, *matplotlib.rcParams["font.sans-serif"]]
matplotlib.rcParams["axes.unicode_minus"] = False # 中文字体里没有 U+2212
return True
logger.warning("未找到中文字体(尝试过 %s),图上改用英文标签", "/".join(CJK_FONTS[:3]))
return False
def save_scatter(path: Path, result, cfg) -> None:
"""把聚类结果画成二维散点 + 质心距离矩阵,放在 preview 拼图的同级目录。
投影用 PCA 而不是 t-SNE/UMAP:后两者为了把簇拉开会扭曲簇间距离,
而这张图的目的就是看簇之间离多远。图上标注了两个主成分的方差占比——
占比低说明高维结构被压掉得多,图上的远近只能参考,以右边的距离矩阵为准。
"""
import matplotlib
matplotlib.use("Agg") # 无显示环境
import matplotlib.pyplot as plt
cjk = setup_cjk_font()
scatter_cfg = dict(cfg.get("scatter", {}))
max_points = int(scatter_cfg.get("max_points", 5000))
points, centroids, ratio = project_2d(result.points, result.centroids)
labels = result.labels
# 点太多时随机抽样:既避免糊成一团,也控制文件大小
if max_points and points.shape[0] > max_points:
generator = torch.Generator().manual_seed(int(cfg.get("seed", 42)))
index = torch.randperm(points.shape[0], generator=generator)[:max_points]
points, labels = points[index], labels[index]
subtitle = f"随机抽样 {max_points} / {result.labels.shape[0]} 点"
subtitle_en = f"sampled {max_points} / {result.labels.shape[0]}"
else:
subtitle = f"全部 {points.shape[0]} 点"
subtitle_en = f"all {points.shape[0]} points"
points = points.cpu().numpy()
centroids_2d = centroids.cpu().numpy()
labels = labels.cpu().numpy()
num_clusters = result.centroids.shape[0]
colors = plt.get_cmap("tab10")(range(num_clusters))
figsize = scatter_cfg.get("figsize", [14, 6])
fig, (ax, ax_matrix) = plt.subplots(
1, 2, figsize=tuple(figsize), gridspec_kw={"width_ratios": [1.6, 1]}
)
for cluster in range(num_clusters):
mask = labels == cluster
ax.scatter(
points[mask, 0],
points[mask, 1],
s=float(scatter_cfg.get("point_size", 6)),
alpha=float(scatter_cfg.get("alpha", 0.35)),
color=colors[cluster],
linewidths=0,
label=f"cluster_{cluster:02d} n={int(result.sizes[cluster])}",
)
ax.scatter(
centroids_2d[:, 0],
centroids_2d[:, 1],
s=220,
marker="X",
c=colors,
edgecolors="black",
linewidths=1.5,
zorder=5,
)
for cluster in range(num_clusters):
ax.annotate(
f"{cluster:02d}",
centroids_2d[cluster],
textcoords="offset points",
xytext=(0, 12),
fontsize=11,
fontweight="bold",
ha="center",
zorder=6,
color="black",
bbox={"boxstyle": "round,pad=0.2", "fc": "white", "ec": "none", "alpha": 0.8},
)
if cjk:
scatter_title = (
f"PCA 二维投影({subtitle})\n"
f"两个主成分方差占比 {ratio[0]:.1%} + {ratio[1]:.1%} = {sum(ratio):.1%}"
)
matrix_title = "质心两两距离(聚类空间,非投影后)"
else:
scatter_title = (
f"PCA 2D projection ({subtitle_en})\n"
f"explained variance {ratio[0]:.1%} + {ratio[1]:.1%} = {sum(ratio):.1%}"
)
matrix_title = "pairwise centroid distance (in cluster space)"
ax.set_title(scatter_title)
ax.set_xlabel("PC1")
ax.set_ylabel("PC2")
ax.legend(loc="best", fontsize=8, framealpha=0.9)
ax.grid(alpha=0.2)
ax.set_aspect("equal", adjustable="datalim")
distances = centroid_distances(result.centroids).cpu().numpy()
image = ax_matrix.imshow(distances, cmap="viridis")
ax_matrix.set_title(matrix_title)
ax_matrix.set_xticks(range(num_clusters))
ax_matrix.set_yticks(range(num_clusters))
ax_matrix.set_xticklabels([f"{i:02d}" for i in range(num_clusters)])
ax_matrix.set_yticklabels([f"{i:02d}" for i in range(num_clusters)])
for i in range(num_clusters):
for j in range(num_clusters):
ax_matrix.text(
j, i, f"{distances[i, j]:.2f}",
ha="center", va="center",
color="white" if distances[i, j] < distances.max() * 0.6 else "black",
fontsize=9,
)
fig.colorbar(image, ax=ax_matrix, fraction=0.046)
fig.tight_layout()
fig.savefig(path, dpi=int(scatter_cfg.get("dpi", 150)))
plt.close(fig)
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("--device", default=None, help="覆盖配置中的 device:auto / cuda / mps / cpu")
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)打印指标表,不产出分组"
)
parser.add_argument(
"--train-only",
action="store_true",
help="只做 SimCLR 自监督训练,存完权重就退出,不抽特征也不聚类",
)
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(args.device or cfg.get("device", "auto"))
except DeviceNotAvailableError as exc:
logger.error("%s", exc)
return 1
method = str(cfg.get("method", "pretrained"))
if args.train_only:
# 这三种组合下 --train-only 要么无事可做、要么与另一个"提前结束"的开关打架,
# 与其跑完才发现什么都没产出,不如立刻说清楚
if method != "simclr":
logger.error("--train-only 只对 method=simclr 有意义,当前 method=%s(没有要训练的东西)", method)
return 1
if cfg.get("simclr", {}).get("checkpoint"):
logger.error(
"--train-only 与 simclr.checkpoint 冲突:checkpoint 会跳过训练,"
"加上 --train-only 就什么都不会做。想接着上次练请改用 simclr.resume。"
)
return 1
if args.sweep_k:
logger.error("--train-only 与 --sweep-k 不能同时用:前者不聚类,后者要靠聚类才能扫")
return 1
dataset = build_dataset(dict(cfg["dataset"]))
logger.info(
"扫描到 %d 张图,来自 %d 个目录;簇数 %d,method=%s,设备 %s",
len(dataset),
len({dataset.source_dir(path) for path in dataset.images}),
num_clusters,
method,
describe_device(device),
)
encoder = build_encoder(cfg, device, work_dir)
if args.train_only:
weights = work_dir / "simclr_backbone.pth"
logger.info("--train-only:训练完成,未做聚类。权重在 %s", weights)
logger.info(
"接着聚类:uv run python scripts/cluster_unlabeled.py --config %s "
"--work-dir <新目录> --set simclr.checkpoint=%s",
cfg.get("config_path", "<配置>"),
weights,
)
return 0
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())