Shaun ZhouandClaude Opus 5
92c88b785c
用户反馈"聚类配置里改不了 lr/weight_decay"，实测是两件事：
1. --set 的键打错字母会凭空造出一个谁也不读的键，训练照跑但改动无效。
实例：simclr.optmizer.weight_decay=0.0（optimizer 少个 i）被静默接受。
现在 apply_overrides 默认只允许覆盖已存在的键，报错时给出同级可用键
与最接近的拼写（"是不是想写 'optimizer'？"）。需要新增键时可 strict=False。
2. lr 其实生效了，但从 2229/2500 续训时 cosine 只剩满值的 3%，
新设的 1e-3 落到 2.87e-5。新增 restart_schedule 开关：按剩余轮数重开曲线，
新 lr 从满值开始；SimCLR 与检测训练都支持。续训日志补上调度进度，
并在 lr 不足 initial_lr 的 10% 时警告并给出两种解法。
顺带：严格校验立刻暴露出 configs/cluster/film.yaml 里根本没有 scatter 段
（代码 cfg.get("scatter", {}) 兜底走默认值，所以功能正常但配置项形同虚设），已补上。
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
726 lines
29 KiB
Python
726 lines
29 KiB
Python
"""SimCLR 自监督：损失的数学性质 + 聚类脚本的 simclr 分支端到端。"""
|
||
|
||
import sys
|
||
|
||
import pytest
|
||
|
||
torch = pytest.importorskip("torch")
|
||
pytest.importorskip("PIL")
|
||
|
||
import numpy as np # noqa: E402
|
||
from PIL import Image # noqa: E402
|
||
|
||
from engine.simclr import SimCLRTrainer, nt_xent_loss # noqa: E402
|
||
from models.heads.projection import ProjectionHead # noqa: E402
|
||
from models.registry import BACKBONES # noqa: E402
|
||
|
||
|
||
def test_nt_xent_is_minimal_when_views_agree():
|
||
"""两个视图完全一致且各样本互相正交时，损失应接近理论下界 log(2B−1) 的反面——
|
||
即正样本对独占概率质量，损失趋近 0。"""
|
||
z = torch.eye(4) # 4 个互相正交的样本
|
||
aligned = nt_xent_loss(z, z.clone(), temperature=0.05)
|
||
misaligned = nt_xent_loss(z, z.roll(1, dims=0), temperature=0.05)
|
||
|
||
assert aligned < 0.1
|
||
assert misaligned > aligned
|
||
|
||
|
||
def test_nt_xent_upper_bound_at_chance_level():
|
||
"""所有表征相同时无法区分正负样本，损失 = log(2B−1)。"""
|
||
z = torch.ones(8, 16)
|
||
loss = float(nt_xent_loss(z, z.clone(), temperature=0.2))
|
||
assert loss == pytest.approx(float(torch.log(torch.tensor(2 * 8 - 1))), abs=1e-4)
|
||
|
||
|
||
def test_nt_xent_is_symmetric_and_temperature_sensitive():
|
||
torch.manual_seed(0)
|
||
z1, z2 = torch.randn(6, 8), torch.randn(6, 8)
|
||
assert float(nt_xent_loss(z1, z2, 0.2)) == pytest.approx(float(nt_xent_loss(z2, z1, 0.2)), abs=1e-5)
|
||
# 温度越低，随机表征的损失越大（softmax 更尖，惩罚更重）
|
||
assert nt_xent_loss(z1, z2, 0.05) > nt_xent_loss(z1, z2, 0.5)
|
||
|
||
|
||
def test_nt_xent_gradient_pulls_positive_pair_together():
|
||
torch.manual_seed(0)
|
||
z1 = torch.randn(4, 8, requires_grad=True)
|
||
z2 = torch.randn(4, 8, requires_grad=True)
|
||
before = float(torch.nn.functional.cosine_similarity(z1, z2).mean().detach())
|
||
|
||
optimizer = torch.optim.SGD([z1, z2], lr=1.0)
|
||
for _ in range(20):
|
||
optimizer.zero_grad()
|
||
nt_xent_loss(z1, z2, 0.2).backward()
|
||
optimizer.step()
|
||
|
||
after = float(torch.nn.functional.cosine_similarity(z1, z2).mean().detach())
|
||
assert after > before
|
||
|
||
|
||
def test_projection_head_accepts_feature_map_and_vector():
|
||
head = ProjectionHead(16, hidden_channels=32, out_channels=8)
|
||
assert head(torch.randn(4, 16, 5, 5)).shape == (4, 8)
|
||
assert head(torch.randn(4, 16)).shape == (4, 8)
|
||
|
||
|
||
@pytest.fixture
|
||
def tiny_image_dir(tmp_path):
|
||
root = tmp_path / "imgs"
|
||
root.mkdir()
|
||
rng = np.random.default_rng(0)
|
||
for i in range(8):
|
||
Image.fromarray(rng.integers(0, 255, (40, 60, 3), dtype=np.uint8)).save(root / f"{i}.png")
|
||
return root
|
||
|
||
|
||
def test_simclr_trainer_runs_and_saves_backbone_only(tiny_image_dir, tmp_path):
|
||
from torch.utils.data import DataLoader
|
||
|
||
from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate
|
||
|
||
dataset = UnlabeledImageDataset(
|
||
tiny_image_dir,
|
||
transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
|
||
views=2,
|
||
view_transforms=[{"name": "RandomResizedCrop", "size": [32, 32], "scale": [0.5, 1.0]}],
|
||
)
|
||
loader = DataLoader(dataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=unlabeled_collate)
|
||
backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})
|
||
|
||
trainer = SimCLRTrainer(
|
||
backbone,
|
||
loader,
|
||
{"epochs": 2, "amp": False, "log_interval": 100, "projection": {"hidden_channels": 16,
|
||
"out_channels": 8}},
|
||
torch.device("cpu"),
|
||
tmp_path / "run",
|
||
)
|
||
path = trainer.fit()
|
||
|
||
assert path.exists()
|
||
state = torch.load(path, map_location="cpu", weights_only=False)
|
||
# 只存骨架，不含投影头
|
||
assert all("projector" not in key and "mlp" not in key for key in state["model"])
|
||
backbone.load_state_dict(state["model"]) # 能被干净地加载回去
|
||
|
||
|
||
def test_simclr_logs_metrics_to_swanlab(tiny_image_dir, tmp_path, monkeypatch):
|
||
"""SwanLab 接线：按 step 记 loss/lr，按 epoch 记平均 loss、耗时与显存峰值，结束时收尾。"""
|
||
from torch.utils.data import DataLoader
|
||
|
||
import engine.simclr as simclr_mod
|
||
from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate
|
||
|
||
records: list[tuple[dict, int | None, str]] = []
|
||
finished: list[bool] = []
|
||
|
||
class StubLogger:
|
||
def __init__(self, **kwargs):
|
||
self.kwargs = kwargs
|
||
|
||
def log(self, metrics, step=None, prefix=""):
|
||
records.append((metrics, step, prefix))
|
||
|
||
def finish(self):
|
||
finished.append(True)
|
||
|
||
monkeypatch.setattr(simclr_mod, "SwanLabLogger", StubLogger)
|
||
|
||
dataset = UnlabeledImageDataset(
|
||
tiny_image_dir,
|
||
transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
|
||
views=2,
|
||
)
|
||
loader = DataLoader(dataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=unlabeled_collate)
|
||
backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})
|
||
|
||
trainer = simclr_mod.SimCLRTrainer(
|
||
backbone,
|
||
loader,
|
||
{"epochs": 1, "amp": False, "log_interval": 1, "save_interval": 0,
|
||
"projection": {"hidden_channels": 16, "out_channels": 8},
|
||
"swanlab": {"enabled": True, "mode": "disabled"}},
|
||
torch.device("cpu"),
|
||
tmp_path / "run",
|
||
full_config={"method": "simclr"},
|
||
)
|
||
trainer.fit()
|
||
|
||
assert records, "没有任何指标被记录"
|
||
assert all(prefix == "simclr/" for _, _, prefix in records)
|
||
step_keys = {k for metrics, _, _ in records for k in metrics}
|
||
assert {"loss", "loss_avg", "lr"} <= step_keys
|
||
assert {"epoch", "epoch_loss", "epoch_seconds", "gpu_mem_peak_mb"} <= step_keys
|
||
assert all(step is not None and step > 0 for _, step, _ in records)
|
||
assert finished == [True] # fit() 结束时调用了 finish()
|
||
|
||
# 完整 yaml 配置作为超参传给了 SwanLab
|
||
assert trainer.visualizer.kwargs["config"] == {"method": "simclr"}
|
||
assert trainer.visualizer.kwargs["mode"] == "disabled"
|
||
|
||
|
||
def test_simclr_saves_intermediate_checkpoints(tiny_image_dir, tmp_path):
|
||
"""中途落盘：跑到一半被打断也要能拿到权重（此前只在全部 epoch 跑完后才存）。"""
|
||
from torch.utils.data import DataLoader
|
||
|
||
from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate
|
||
|
||
dataset = UnlabeledImageDataset(
|
||
tiny_image_dir,
|
||
transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
|
||
views=2,
|
||
)
|
||
loader = DataLoader(dataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=unlabeled_collate)
|
||
backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})
|
||
work_dir = tmp_path / "run"
|
||
|
||
trainer = SimCLRTrainer(
|
||
backbone,
|
||
loader,
|
||
{"epochs": 4, "amp": False, "log_interval": 100, "save_interval": 2,
|
||
"projection": {"hidden_channels": 16, "out_channels": 8}},
|
||
torch.device("cpu"),
|
||
work_dir,
|
||
)
|
||
# 手动走前两个 epoch，模拟"训练中途"
|
||
trainer.train_one_epoch(0)
|
||
trainer.train_one_epoch(1)
|
||
path = trainer.save_backbone(1)
|
||
|
||
assert path.exists()
|
||
assert torch.load(path, map_location="cpu", weights_only=False)["epochs"] == 2
|
||
|
||
trainer.fit() # 完整跑完后 epochs 应更新
|
||
assert torch.load(path, map_location="cpu", weights_only=False)["epochs"] == 4
|
||
|
||
|
||
def _make_trainer(tiny_image_dir, work_dir, epochs=4, **cfg_extra):
|
||
from torch.utils.data import DataLoader
|
||
|
||
from datasets.unlabeled import UnlabeledImageDataset, unlabeled_collate
|
||
|
||
dataset = UnlabeledImageDataset(
|
||
tiny_image_dir,
|
||
transforms=[{"name": "Resize", "size": [32, 32]}, {"name": "ToTensor"}],
|
||
views=2,
|
||
)
|
||
loader = DataLoader(dataset, batch_size=4, shuffle=False, drop_last=True, collate_fn=unlabeled_collate)
|
||
torch.manual_seed(0)
|
||
backbone = BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]})
|
||
cfg = {
|
||
"epochs": epochs,
|
||
"amp": False,
|
||
"log_interval": 100,
|
||
"save_interval": 1,
|
||
"projection": {"hidden_channels": 16, "out_channels": 8},
|
||
"scheduler": {"name": "cosine"},
|
||
"swanlab": {"enabled": False},
|
||
**cfg_extra,
|
||
}
|
||
return SimCLRTrainer(backbone, loader, cfg, torch.device("cpu"), work_dir,
|
||
resume=cfg.pop("resume", None))
|
||
|
||
|
||
def test_save_state_contains_full_training_state(tiny_image_dir, tmp_path):
|
||
trainer = _make_trainer(tiny_image_dir, tmp_path / "run", epochs=1)
|
||
trainer.fit()
|
||
|
||
state_path = tmp_path / "run" / "simclr_last.pth"
|
||
assert state_path.exists()
|
||
state = torch.load(state_path, map_location="cpu", weights_only=False)
|
||
# 续训需要的东西一个都不能少，否则动量和学习率进度会被清零
|
||
assert {"backbone", "projector", "optimizer", "scheduler", "scaler", "epoch", "global_step"} <= set(state)
|
||
assert state["epoch"] == 0
|
||
|
||
|
||
def test_resume_continues_from_the_saved_epoch(tiny_image_dir, tmp_path):
|
||
"""resume 是接着练，不是从头再来。"""
|
||
work_dir = tmp_path / "run"
|
||
first = _make_trainer(tiny_image_dir, work_dir, epochs=2)
|
||
first.fit()
|
||
steps_after_2 = first.global_step
|
||
|
||
resumed = SimCLRTrainer(
|
||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||
first.dataloader,
|
||
{"epochs": 5, "amp": False, "log_interval": 100, "save_interval": 1,
|
||
"projection": {"hidden_channels": 16, "out_channels": 8},
|
||
"scheduler": {"name": "cosine"}, "swanlab": {"enabled": False}},
|
||
torch.device("cpu"),
|
||
work_dir,
|
||
resume=work_dir / "simclr_last.pth",
|
||
)
|
||
assert resumed.start_epoch == 2 # 从第 3 轮开始
|
||
assert resumed.global_step == steps_after_2 # step 计数接上
|
||
# 总轮数从 2 改成了 5，学习率按新计划重算（若沿用旧的 T_max，此处会是 0）
|
||
assert resumed.optimizer.param_groups[0]["lr"] > 1e-5
|
||
|
||
resumed.fit()
|
||
state = torch.load(work_dir / "simclr_last.pth", map_location="cpu", weights_only=False)
|
||
assert state["epoch"] == 4 # 一路练到第 5 轮
|
||
assert resumed.global_step > steps_after_2
|
||
|
||
|
||
def test_resume_with_a_longer_schedule_rebuilds_the_lr_curve(tiny_image_dir, tmp_path):
|
||
"""续训时调大 epochs，学习率必须按新的总轮数重算。
|
||
|
||
回归测试：此前直接恢复调度器状态，会沿用存档里的 T_max。原计划 10 轮的 cosine
|
||
跑完后 lr 已经是 0，这时把 epochs 改成 30 接着练，lr 仍是 0——训练照跑但一步都学不动。
|
||
"""
|
||
import math
|
||
|
||
work_dir = tmp_path / "run"
|
||
first = _make_trainer(tiny_image_dir, work_dir, epochs=10)
|
||
first.fit()
|
||
assert first.optimizer.param_groups[0]["lr"] == pytest.approx(0.0, abs=1e-9) # cosine 已走完
|
||
|
||
resumed = SimCLRTrainer(
|
||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||
first.dataloader,
|
||
{"epochs": 30, "amp": False, "log_interval": 999, "save_interval": 1,
|
||
"scheduler": {"name": "cosine"}, "optimizer": {"name": "adamw", "lr": 1e-3},
|
||
"projection": {"hidden_channels": 16, "out_channels": 8},
|
||
"swanlab": {"enabled": False}},
|
||
torch.device("cpu"),
|
||
work_dir,
|
||
resume=work_dir / "simclr_last.pth",
|
||
)
|
||
|
||
assert resumed.scheduler.T_max == 30
|
||
lr = resumed.optimizer.param_groups[0]["lr"]
|
||
assert lr > 1e-5, "续训后学习率仍是 0，等于白跑"
|
||
assert lr == pytest.approx(1e-3 * (1 + math.cos(math.pi * 10 / 30)) / 2, rel=1e-3)
|
||
|
||
|
||
def test_restart_schedule_gives_the_new_lr_full_value(tiny_image_dir, tmp_path):
|
||
"""续训到调度末段时，cosine 会把新 lr 乘上接近 0 的系数；
|
||
restart_schedule 按剩余轮数重开曲线，让新 lr 从满值开始。"""
|
||
import math
|
||
|
||
work_dir = tmp_path / "run"
|
||
_make_trainer(
|
||
tiny_image_dir, work_dir, epochs=9,
|
||
optimizer={"name": "adamw", "lr": 1e-3},
|
||
).fit()
|
||
|
||
def resumed(restart):
|
||
return SimCLRTrainer(
|
||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||
_make_trainer(tiny_image_dir, work_dir, epochs=1).dataloader,
|
||
{"epochs": 10, "amp": False, "log_interval": 999, "save_interval": 1,
|
||
"scheduler": {"name": "cosine"}, "optimizer": {"name": "adamw", "lr": 1e-3},
|
||
"projection": {"hidden_channels": 16, "out_channels": 8},
|
||
"swanlab": {"enabled": False}, "restart_schedule": restart},
|
||
torch.device("cpu"),
|
||
work_dir,
|
||
resume=work_dir / "simclr_last.pth",
|
||
)
|
||
|
||
# 默认：沿用整条曲线，9/10 处只剩很小的系数
|
||
default_lr = resumed(False).optimizer.param_groups[0]["lr"]
|
||
assert default_lr == pytest.approx(1e-3 * (1 + math.cos(math.pi * 9 / 10)) / 2)
|
||
assert default_lr < 1e-3 * 0.1
|
||
|
||
# 打开开关：新 lr 从满值开始
|
||
assert resumed(True).optimizer.param_groups[0]["lr"] == pytest.approx(1e-3)
|
||
|
||
|
||
def test_resume_takes_hyperparams_from_config_not_checkpoint(tiny_image_dir, tmp_path):
|
||
"""续训时 lr / weight_decay / 总轮数一律以 yaml 为准，存档只提供进度。
|
||
|
||
回归测试：optimizer.load_state_dict 会整体覆盖 param_groups，把超参连同动量缓冲
|
||
一起还原成存档时的值——改了配置再续训，改动被静默吃掉。
|
||
"""
|
||
import math
|
||
|
||
work_dir = tmp_path / "run"
|
||
first = _make_trainer(
|
||
tiny_image_dir, work_dir, epochs=4,
|
||
optimizer={"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4},
|
||
)
|
||
first.fit()
|
||
|
||
resumed = SimCLRTrainer(
|
||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||
first.dataloader,
|
||
{"epochs": 8, "amp": False, "log_interval": 999, "save_interval": 1,
|
||
"scheduler": {"name": "cosine"},
|
||
"optimizer": {"name": "adamw", "lr": 5e-4, "weight_decay": 0.05},
|
||
"projection": {"hidden_channels": 16, "out_channels": 8},
|
||
"swanlab": {"enabled": False}},
|
||
torch.device("cpu"),
|
||
work_dir,
|
||
resume=work_dir / "simclr_last.pth",
|
||
)
|
||
|
||
groups = resumed.optimizer.param_groups
|
||
assert groups[0]["initial_lr"] == pytest.approx(5e-4) # 不是存档里的 1e-3
|
||
assert groups[0]["weight_decay"] == pytest.approx(0.05) # 不是存档里的 1e-4
|
||
assert groups[1]["weight_decay"] == 0.0 # BN/bias 分组结构没被冲掉
|
||
assert resumed.scheduler.base_lrs == pytest.approx([5e-4, 5e-4])
|
||
assert resumed.scheduler.T_max == 8
|
||
assert groups[0]["lr"] == pytest.approx(5e-4 * (1 + math.cos(math.pi * 4 / 8)) / 2)
|
||
assert resumed.optimizer.state, "动量缓冲应当保留——那才是真正要续的'进度'"
|
||
|
||
|
||
def test_resume_across_optimizer_types_drops_incompatible_state(tiny_image_dir, tmp_path):
|
||
"""换优化器种类时不加载动量缓冲：AdamW 的 exp_avg 对 SGD 没有意义，硬塞只是白占内存。"""
|
||
work_dir = tmp_path / "run"
|
||
_make_trainer(
|
||
tiny_image_dir, work_dir, epochs=2,
|
||
optimizer={"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4},
|
||
).fit()
|
||
|
||
resumed = SimCLRTrainer(
|
||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||
_make_trainer(tiny_image_dir, work_dir, epochs=1).dataloader,
|
||
{"epochs": 6, "amp": False, "log_interval": 999, "save_interval": 1,
|
||
"scheduler": {"name": "cosine"},
|
||
"optimizer": {"name": "sgd", "lr": 0.01, "momentum": 0.9},
|
||
"projection": {"hidden_channels": 16, "out_channels": 8},
|
||
"swanlab": {"enabled": False}},
|
||
torch.device("cpu"),
|
||
work_dir,
|
||
resume=work_dir / "simclr_last.pth",
|
||
)
|
||
|
||
assert isinstance(resumed.optimizer, torch.optim.SGD)
|
||
assert not resumed.optimizer.state, "不该带着 AdamW 的缓冲"
|
||
assert resumed.start_epoch == 2 # 训练进度仍然接上
|
||
assert resumed.optimizer.param_groups[0]["momentum"] == 0.9
|
||
|
||
|
||
def test_resume_with_same_schedule_restores_state(tiny_image_dir, tmp_path):
|
||
"""总轮数不变时，重建 + 快进的结果应与原调度进度一致。"""
|
||
work_dir = tmp_path / "run"
|
||
_make_trainer(tiny_image_dir, work_dir, epochs=6).fit()
|
||
|
||
resumed = SimCLRTrainer(
|
||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||
_make_trainer(tiny_image_dir, work_dir, epochs=1).dataloader,
|
||
{"epochs": 6, "amp": False, "log_interval": 999, "save_interval": 1,
|
||
"scheduler": {"name": "cosine"},
|
||
"projection": {"hidden_channels": 16, "out_channels": 8},
|
||
"swanlab": {"enabled": False}},
|
||
torch.device("cpu"),
|
||
work_dir,
|
||
resume=work_dir / "simclr_last.pth",
|
||
)
|
||
assert resumed.scheduler.T_max == 6
|
||
assert resumed.scheduler.last_epoch == 6
|
||
|
||
|
||
def test_intermediate_checkpoint_is_resumable(tiny_image_dir, tmp_path):
|
||
"""save_interval 落的中间盘必须能直接拿来续训（不必等训练正常结束）。"""
|
||
work_dir = tmp_path / "run"
|
||
trainer = _make_trainer(tiny_image_dir, work_dir, epochs=10)
|
||
# 只手动跑到第 4 轮就落盘，模拟被 kill 掉
|
||
for epoch in range(4):
|
||
trainer.train_one_epoch(epoch)
|
||
trainer.scheduler.step()
|
||
trainer._checkpoint(3)
|
||
|
||
resumed = SimCLRTrainer(
|
||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||
trainer.dataloader,
|
||
{"epochs": 10, "amp": False, "log_interval": 999, "save_interval": 1,
|
||
"scheduler": {"name": "cosine"},
|
||
"projection": {"hidden_channels": 16, "out_channels": 8},
|
||
"swanlab": {"enabled": False}},
|
||
torch.device("cpu"),
|
||
work_dir,
|
||
resume=work_dir / "simclr_last.pth",
|
||
)
|
||
assert resumed.start_epoch == 4
|
||
resumed.fit()
|
||
state = torch.load(work_dir / "simclr_last.pth", map_location="cpu", weights_only=False)
|
||
assert state["epoch"] == 9
|
||
|
||
|
||
def test_resume_rejects_backbone_only_file(tiny_image_dir, tmp_path):
|
||
"""simclr_backbone.pth 只有骨架，不能当续训状态用，要给出可操作的报错。"""
|
||
work_dir = tmp_path / "run"
|
||
_make_trainer(tiny_image_dir, work_dir, epochs=1).fit()
|
||
|
||
with pytest.raises(ValueError, match="只含骨架权重"):
|
||
SimCLRTrainer(
|
||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||
_make_trainer(tiny_image_dir, work_dir, epochs=1).dataloader,
|
||
{"epochs": 2, "amp": False, "projection": {"hidden_channels": 16, "out_channels": 8},
|
||
"swanlab": {"enabled": False}},
|
||
torch.device("cpu"),
|
||
work_dir,
|
||
resume=work_dir / "simclr_backbone.pth",
|
||
)
|
||
|
||
with pytest.raises(FileNotFoundError, match="不存在"):
|
||
_make_trainer(tiny_image_dir, work_dir, epochs=1, resume=work_dir / "nope.pth")
|
||
|
||
|
||
def test_resume_beyond_planned_epochs_does_not_train(tiny_image_dir, tmp_path):
|
||
work_dir = tmp_path / "run"
|
||
_make_trainer(tiny_image_dir, work_dir, epochs=3).fit()
|
||
|
||
resumed = SimCLRTrainer(
|
||
BACKBONES.build({"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}),
|
||
_make_trainer(tiny_image_dir, work_dir, epochs=1).dataloader,
|
||
{"epochs": 3, "amp": False, "projection": {"hidden_channels": 16, "out_channels": 8},
|
||
"swanlab": {"enabled": False}},
|
||
torch.device("cpu"),
|
||
work_dir,
|
||
resume=work_dir / "simclr_last.pth",
|
||
)
|
||
steps_before = resumed.global_step
|
||
resumed.fit()
|
||
assert resumed.global_step == steps_before # 一步都没再训
|
||
|
||
|
||
def test_resume_and_checkpoint_are_mutually_exclusive(tiny_image_dir, tmp_path, monkeypatch):
|
||
sys.path.insert(0, "scripts")
|
||
import cluster_unlabeled as cu
|
||
|
||
with pytest.raises(ValueError, match="只能二选一"):
|
||
cu.build_encoder(
|
||
{"method": "simclr",
|
||
"simclr": {"resume": "a.pth", "checkpoint": "b.pth"},
|
||
"dataset": {"name": "unlabeled", "root": str(tiny_image_dir)}},
|
||
torch.device("cpu"),
|
||
tmp_path,
|
||
)
|
||
|
||
|
||
def test_cluster_script_simclr_branch_end_to_end(tiny_image_dir, tmp_path, monkeypatch):
|
||
"""method=simclr 全流程：训练 -> 存权重 -> 抽特征 -> 聚类 -> 产出。"""
|
||
sys.path.insert(0, "scripts")
|
||
import cluster_unlabeled as cu
|
||
|
||
work_dir = tmp_path / "run"
|
||
monkeypatch.setattr(
|
||
sys, "argv",
|
||
["cluster_unlabeled.py", "--config", "configs/cluster/film.yaml",
|
||
"--work-dir", str(work_dir), "--root", str(tiny_image_dir), "--set",
|
||
"device=cpu", "method=simclr", "loader.num_workers=0", "loader.batch_size=4",
|
||
"simclr.epochs=1", "simclr.batch_size=4", "simclr.amp=False",
|
||
"simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]",
|
||
"simclr.projection.hidden_channels=16", "simclr.projection.out_channels=8",
|
||
"num_clusters=2", "n_init=2"],
|
||
)
|
||
|
||
assert cu.main() == 0
|
||
assert (work_dir / "simclr_backbone.pth").exists()
|
||
assert (work_dir / "assignments.csv").exists()
|
||
assert (work_dir / "summary.json").exists()
|
||
assert (work_dir / "clusters" / "cluster_00").is_dir()
|
||
|
||
|
||
def _cluster_argv(work_dir, image_dir, *extra):
|
||
return [
|
||
"cluster_unlabeled.py", "--config", "configs/cluster/film.yaml",
|
||
"--work-dir", str(work_dir), "--root", str(image_dir), "--set",
|
||
"device=cpu", "loader.num_workers=0", "loader.batch_size=4",
|
||
"num_clusters=2", "n_init=2", *extra,
|
||
]
|
||
|
||
|
||
def test_warns_when_simclr_keys_are_at_top_level(tiny_image_dir, caplog):
|
||
"""顶层写 optimizer 会被静默忽略，必须警告——否则"改了配置没反应"极难排查。"""
|
||
import logging
|
||
|
||
sys.path.insert(0, "scripts")
|
||
import cluster_unlabeled as cu
|
||
|
||
cfg = {
|
||
"dataset": {"name": "unlabeled", "root": str(tiny_image_dir)},
|
||
"optimizer": {"name": "adamw", "lr": 0.9}, # ← 应该放在 simclr: 下
|
||
"epochs": 100, # ← 同上
|
||
"simclr": {"epochs": 2},
|
||
}
|
||
with caplog.at_level(logging.WARNING, logger="mte"):
|
||
cu.warn_misplaced_keys(cfg, "simclr")
|
||
|
||
assert "optimizer" in caplog.text and "epochs" in caplog.text
|
||
assert "simclr" in caplog.text
|
||
|
||
|
||
def test_no_warning_when_keys_are_correctly_nested(tiny_image_dir, caplog):
|
||
import logging
|
||
|
||
sys.path.insert(0, "scripts")
|
||
import cluster_unlabeled as cu
|
||
|
||
cfg = {"dataset": {"name": "unlabeled", "root": str(tiny_image_dir)},
|
||
"simclr": {"optimizer": {"lr": 0.9}, "epochs": 2}}
|
||
with caplog.at_level(logging.WARNING, logger="mte"):
|
||
cu.warn_misplaced_keys(cfg, "simclr")
|
||
assert "会被忽略" not in caplog.text
|
||
|
||
|
||
def test_effective_hyperparams_are_logged_on_every_run(tiny_image_dir, tmp_path, caplog):
|
||
"""不只是续训时——每次训练开始都要打印实际生效的超参。"""
|
||
import logging
|
||
|
||
trainer = _make_trainer(
|
||
tiny_image_dir, tmp_path / "run", epochs=1,
|
||
optimizer={"name": "adamw", "lr": 0.00042, "weight_decay": 0.007},
|
||
)
|
||
del trainer
|
||
with caplog.at_level(logging.INFO, logger="mte"):
|
||
_make_trainer(
|
||
tiny_image_dir, tmp_path / "run2", epochs=1,
|
||
optimizer={"name": "adamw", "lr": 0.00042, "weight_decay": 0.007},
|
||
)
|
||
assert "SimCLR 超参" in caplog.text
|
||
assert "0.00042" in caplog.text
|
||
|
||
|
||
def test_train_only_stops_before_clustering(tiny_image_dir, tmp_path, monkeypatch):
|
||
"""--train-only 只训练存权重，不抽特征也不聚类。"""
|
||
sys.path.insert(0, "scripts")
|
||
import cluster_unlabeled as cu
|
||
|
||
work_dir = tmp_path / "train_only"
|
||
monkeypatch.setattr(
|
||
sys, "argv",
|
||
[*_cluster_argv(work_dir, tiny_image_dir,
|
||
"method=simclr", "simclr.epochs=1", "simclr.batch_size=4",
|
||
"simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]",
|
||
"simclr.projection.hidden_channels=16", "simclr.projection.out_channels=8"),
|
||
"--train-only"],
|
||
)
|
||
assert cu.main() == 0
|
||
|
||
assert (work_dir / "simclr_backbone.pth").exists()
|
||
assert (work_dir / "simclr_last.pth").exists() # 续训状态照常落盘
|
||
for produced_by_clustering in ("assignments.csv", "summary.json", "cluster_scatter.png", "clusters"):
|
||
assert not (work_dir / produced_by_clustering).exists(), produced_by_clustering
|
||
|
||
|
||
def test_train_only_then_cluster_only(tiny_image_dir, tmp_path, monkeypatch):
|
||
"""训练与聚类可以分两次跑：第一次只训练，第二次用 checkpoint 只聚类。"""
|
||
sys.path.insert(0, "scripts")
|
||
import cluster_unlabeled as cu
|
||
|
||
train_dir = tmp_path / "phase1"
|
||
monkeypatch.setattr(
|
||
sys, "argv",
|
||
[*_cluster_argv(train_dir, tiny_image_dir,
|
||
"method=simclr", "simclr.epochs=1", "simclr.batch_size=4",
|
||
"simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]",
|
||
"simclr.projection.hidden_channels=16", "simclr.projection.out_channels=8"),
|
||
"--train-only"],
|
||
)
|
||
assert cu.main() == 0
|
||
|
||
cluster_dir = tmp_path / "phase2"
|
||
monkeypatch.setattr(
|
||
sys, "argv",
|
||
_cluster_argv(cluster_dir, tiny_image_dir,
|
||
"method=simclr",
|
||
f"simclr.checkpoint={train_dir / 'simclr_backbone.pth'}",
|
||
"simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]"),
|
||
)
|
||
assert cu.main() == 0
|
||
|
||
assert (cluster_dir / "assignments.csv").exists()
|
||
assert not (cluster_dir / "simclr_backbone.pth").exists(), "第二步不该重新训练"
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"extra,reason",
|
||
[
|
||
([], "method=pretrained 没有要训练的东西"),
|
||
(["method=simclr", "simclr.checkpoint=whatever.pth"], "checkpoint 会跳过训练"),
|
||
],
|
||
)
|
||
def test_train_only_rejects_meaningless_combinations(
|
||
tiny_image_dir, tmp_path, monkeypatch, extra, reason
|
||
):
|
||
"""这些组合下 --train-only 什么都不会做，应当立刻报错而不是跑完才发现没产出。"""
|
||
sys.path.insert(0, "scripts")
|
||
import cluster_unlabeled as cu
|
||
|
||
monkeypatch.setattr(
|
||
sys, "argv", [*_cluster_argv(tmp_path / "bad", tiny_image_dir, *extra), "--train-only"]
|
||
)
|
||
assert cu.main() == 1, reason
|
||
|
||
|
||
def test_train_only_conflicts_with_sweep_k(tiny_image_dir, tmp_path, monkeypatch):
|
||
sys.path.insert(0, "scripts")
|
||
import cluster_unlabeled as cu
|
||
|
||
monkeypatch.setattr(
|
||
sys, "argv",
|
||
[*_cluster_argv(tmp_path / "bad", tiny_image_dir, "method=simclr"),
|
||
"--train-only", "--sweep-k", "2-4"],
|
||
)
|
||
assert cu.main() == 1
|
||
|
||
|
||
def test_scatter_plot_is_written_next_to_previews(tiny_image_dir, tmp_path, monkeypatch):
|
||
"""cluster_scatter.png 要和 preview 拼图落在同一级目录。"""
|
||
sys.path.insert(0, "scripts")
|
||
import cluster_unlabeled as cu
|
||
|
||
work_dir = tmp_path / "run"
|
||
monkeypatch.setattr(
|
||
sys, "argv",
|
||
["cluster_unlabeled.py", "--config", "configs/cluster/film.yaml",
|
||
"--work-dir", str(work_dir), "--root", str(tiny_image_dir), "--set",
|
||
"device=cpu", "loader.num_workers=0", "loader.batch_size=4",
|
||
"num_clusters=2", "n_init=2"],
|
||
)
|
||
assert cu.main() == 0
|
||
|
||
scatter = work_dir / "cluster_scatter.png"
|
||
assert scatter.exists() and scatter.stat().st_size > 5000
|
||
assert (work_dir / "cluster_00_preview.png").parent == scatter.parent
|
||
|
||
import json
|
||
|
||
summary = json.loads((work_dir / "summary.json").read_text(encoding="utf-8"))
|
||
matrix = summary["centroid_distances"]
|
||
assert len(matrix) == 2 and matrix[0][0] == pytest.approx(0.0)
|
||
assert matrix[0][1] == pytest.approx(matrix[1][0])
|
||
|
||
|
||
def test_scatter_can_be_disabled(tiny_image_dir, tmp_path, monkeypatch):
|
||
sys.path.insert(0, "scripts")
|
||
import cluster_unlabeled as cu
|
||
|
||
work_dir = tmp_path / "run_off"
|
||
monkeypatch.setattr(
|
||
sys, "argv",
|
||
["cluster_unlabeled.py", "--config", "configs/cluster/film.yaml",
|
||
"--work-dir", str(work_dir), "--root", str(tiny_image_dir), "--set",
|
||
"device=cpu", "loader.num_workers=0", "loader.batch_size=4",
|
||
"num_clusters=2", "n_init=2", "scatter.enabled=False"],
|
||
)
|
||
assert cu.main() == 0
|
||
assert not (work_dir / "cluster_scatter.png").exists()
|
||
assert (work_dir / "cluster_00_preview.png").exists()
|
||
|
||
|
||
def test_simclr_runs_on_the_resolved_device(tiny_image_dir, tmp_path, monkeypatch):
|
||
"""method=simclr 在本机自动选中的后端上（MPS / CUDA / CPU）都应能跑通。"""
|
||
sys.path.insert(0, "scripts")
|
||
import cluster_unlabeled as cu
|
||
|
||
from utils.device import describe_device, resolve_device
|
||
|
||
device = resolve_device("auto")
|
||
work_dir = tmp_path / "auto_run"
|
||
monkeypatch.setattr(
|
||
sys, "argv",
|
||
["cluster_unlabeled.py", "--config", "configs/cluster/film.yaml",
|
||
"--work-dir", str(work_dir), "--root", str(tiny_image_dir), "--set",
|
||
"method=simclr", "loader.num_workers=0", "loader.batch_size=4",
|
||
"simclr.epochs=1", "simclr.batch_size=4",
|
||
"simclr.backbone.widths=[8,16]", "simclr.backbone.depths=[1,1]",
|
||
"simclr.projection.hidden_channels=16", "simclr.projection.out_channels=8",
|
||
"num_clusters=2", "n_init=2"],
|
||
)
|
||
|
||
assert cu.main() == 0, f"在 {describe_device(device)} 上跑失败"
|
||
assert (work_dir / "simclr_backbone.pth").exists()
|