Shaun ZhouandClaude Opus 5
33d28dcc7c
软链接记的是绝对路径，从服务器拷到本机后全部失效，也不便于直接在 runs/ 下整理。 硬链接不占额外磁盘却表现得和真实文件一样：看图软件能开、可以随意挪动、建子目录 继续分子类，整个目录拷到别的机器也不会断链。 - 新增 link_mode: hard（默认）/ sym / copy；跨文件系统时 hard 自动退回软链并警告 - 落盘后打印各方式的条目数，便于确认实际用的是哪种 - export_clusters.py 适配硬链接：条目本身就是原图，但路径在 runs/ 下、名字还带 0125_ 排序前缀。按 inode 反查 assignments.csv 还原原图路径，这样导出的文件名与 --prefix-parent 才能拿到真正的来源目录；查不到则剥掉前缀后用文件名索引兜底 实测：14,579 张跑聚类后 clusters/ 全为硬链接（nlink=2、可直接打开）； 把簇改名、建子目录挪图后再导出，排序前缀已剥除、来源目录正确还原为「反面纵黑点」， 120 张内容与原图逐字节一致。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
587 lines
23 KiB
Python
587 lines
23 KiB
Python
#!/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 errno
|
||
import json
|
||
import os
|
||
import shutil
|
||
import sys
|
||
from collections import Counter, 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))
|
||
link_mode = str(cfg.get("link_mode", "hard")).lower()
|
||
if link_mode not in LINK_MODES:
|
||
raise ValueError(f"未知 link_mode: {link_mode}，可选 {LINK_MODES}")
|
||
clusters_dir = work_dir / "clusters"
|
||
if clusters_dir.exists():
|
||
shutil.rmtree(clusters_dir)
|
||
|
||
used_modes: Counter[str] = Counter()
|
||
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():
|
||
used_modes[make_link(source, link, link_mode)] += 1
|
||
save_preview(work_dir / f"cluster_{cluster:02d}_preview.png", metas, indices[:preview_count])
|
||
|
||
if used_modes:
|
||
detail = "、".join(f"{mode} {count}" for mode, count in sorted(used_modes.items()))
|
||
logger.info("clusters/ 下的条目：%s", detail)
|
||
if link_mode == "hard" and used_modes.get("sym"):
|
||
logger.warning(
|
||
"有 %d 个条目跨文件系统、退回了软链——这些在别的机器上会断链，"
|
||
"可改用 link_mode: copy，或让 runs/ 与数据同盘",
|
||
used_modes["sym"],
|
||
)
|
||
|
||
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)
|
||
|
||
|
||
#: 这些键只在 `simclr:` 段下有意义，写到顶层会被静默忽略
|
||
SIMCLR_ONLY_KEYS = (
|
||
"optimizer",
|
||
"scheduler",
|
||
"epochs",
|
||
"temperature",
|
||
"projection",
|
||
"save_interval",
|
||
"resume",
|
||
"checkpoint",
|
||
"log_interval",
|
||
)
|
||
|
||
|
||
def warn_misplaced_keys(cfg, method: str) -> None:
|
||
"""顶层出现本该放在 simclr: 下的键时警告——否则改了配置却毫无反应，极难排查。"""
|
||
misplaced = [key for key in SIMCLR_ONLY_KEYS if key in cfg]
|
||
if misplaced:
|
||
logger.warning(
|
||
"配置顶层出现了只在 `simclr:` 段下生效的键 %s，它们会被忽略。"
|
||
"把它们缩进到 simclr: 下面，例如 `simclr:\n optimizer: {lr: 0.0005}`",
|
||
misplaced,
|
||
)
|
||
if method != "simclr" and cfg.get("simclr"):
|
||
logger.info("当前 method=%s，`simclr:` 段下的配置不会生效", method)
|
||
|
||
|
||
LINK_MODES = ("hard", "sym", "copy")
|
||
|
||
|
||
def make_link(source: Path, destination: Path, mode: str) -> str:
|
||
"""在 clusters/ 里放一份指向原图的条目，返回实际用的方式。
|
||
|
||
默认硬链接：不占额外磁盘，但表现得和真实文件一模一样——看图软件、文件管理器、
|
||
以及在 runs/ 下直接继续分子类的工具都能正常处理，把整个目录拷到别的机器上
|
||
也不会变成断链（软链会，因为它记的是绝对路径）。
|
||
|
||
代价是硬链接必须与源文件同一个文件系统；跨盘时自动降级并说明原因。
|
||
另外硬链接和原图是同一份数据，**在 clusters/ 里改图会改到原图**，只挪动、
|
||
删除条目则不影响原图。
|
||
"""
|
||
if mode == "sym":
|
||
destination.symlink_to(source)
|
||
return "sym"
|
||
if mode == "copy":
|
||
shutil.copy2(source, destination)
|
||
return "copy"
|
||
|
||
try:
|
||
os.link(source, destination)
|
||
return "hard"
|
||
except OSError as exc:
|
||
if exc.errno != errno.EXDEV:
|
||
raise
|
||
# 跨文件系统只能退回软链；这时拷到别的机器上会断链，用 export_clusters.py 重定位
|
||
destination.symlink_to(source)
|
||
return "sym"
|
||
|
||
|
||
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
|
||
|
||
warn_misplaced_keys(cfg, method)
|
||
|
||
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())
|