Shaun ZhouandClaude Opus 5
345dee50da
此前 simclr.checkpoint 的语义是"有权重就跳过训练"，没有真正的续训： 中途停掉再启动只能从头再来，或者放弃已训的轮次。 - 新增 simclr_last.pth：骨架 + 投影头 + 优化器 + 调度器 + scaler + epoch + step， 与给下游用的 simclr_backbone.pth 并存，两者在同一时机落盘 - 新增 simclr.resume：载入完整状态从下一个 epoch 接着练。只恢复权重不恢复动量与 cosine 进度的话，续训时损失会先反弹一段 - resume 与 checkpoint 互斥（同时给直接报错）；拿只含骨架的 simclr_backbone.pth 去 resume 会给出可操作的报错，提示改用 checkpoint - 续训轮数已达计划值时不再训练，只打警告 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
307 lines
12 KiB
Python
307 lines
12 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 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 ( # 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])
|
||
|
||
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("--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）打印指标表，不产出分组"
|
||
)
|
||
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
|
||
|
||
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,
|
||
cfg.get("method", "pretrained"),
|
||
describe_device(device),
|
||
)
|
||
|
||
encoder = build_encoder(cfg, device, work_dir)
|
||
|
||
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())
|