Shaun ZhouandClaude Opus 5
105cb38831
train.py --resume 走的是与 SimCLR 不同的一套代码，同样的问题在这里更严重： - load_checkpoint 直接 optimizer.load_state_dict，把 lr/weight_decay 连同动量缓冲 一起还原成存档值：yaml 写 lr=5e-4，实际仍按检查点的 1e-3 训 - save_checkpoint 根本没保存调度器状态，续训时 scheduler.last_epoch 从 0 开始， 而 optimizer 的 lr 又来自检查点，两者互相矛盾，cosine 曲线完全错乱 改为与 SimCLR 一致的「配置文件是唯一真相源」： - save_checkpoint 补上 scheduler 状态、global_step 与当时的 train 配置快照 - 新增 Trainer.resume()：加载优化器状态后按组写回配置里的超参（只留动量缓冲， 并保住 BN/bias 那组 weight_decay=0 的结构），调度器按当前配置重建再快进到已训 epoch - 配置与检查点不一致时逐项打日志说明按当前配置执行 - train.py 改用 trainer.resume()，不再自己拼恢复逻辑 实测续训到 10 轮时 lr=3.273e-04，与 cosine 在 4/10 处的理论值精确一致； weight_decay、T_max、调度进度、动量缓冲均正确。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
218 lines
7.9 KiB
Python
218 lines
7.9 KiB
Python
"""train.py / eval.py 的端到端冒烟：YOLO 数据目录 -> 训练 -> 评估。
|
|
|
|
覆盖入口脚本的接线：classes 解析、label_mode 下发、设备解析。固定用 CPU 跑，
|
|
免得测试受机器上有没有 GPU 影响。
|
|
"""
|
|
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("torch")
|
|
pytest.importorskip("yaml")
|
|
pytest.importorskip("PIL")
|
|
|
|
import numpy as np # noqa: E402
|
|
import torch # noqa: E402
|
|
import yaml # noqa: E402
|
|
from PIL import Image # noqa: E402
|
|
|
|
import eval as eval_mod # noqa: E402
|
|
import train as train_mod # noqa: E402
|
|
from datasets import write_data_yaml # noqa: E402
|
|
|
|
CLASSES = ["scratch", "bubble"]
|
|
TRANSFORMS = [
|
|
{"name": "Resize", "size": [32, 32]},
|
|
{"name": "ToTensor"},
|
|
{"name": "Normalize", "mean": [0.5, 0.5, 0.5], "std": [0.5, 0.5, 0.5]},
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def project(tmp_path):
|
|
"""造一个最小的 YOLO 数据集 + 配置，classes 用 data.yaml 的路径给。"""
|
|
for split in ("train", "val"):
|
|
images = tmp_path / "images" / split
|
|
labels = tmp_path / "labels" / split
|
|
images.mkdir(parents=True)
|
|
labels.mkdir(parents=True)
|
|
for i in range(2):
|
|
Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(images / f"{i}.png")
|
|
content = "0 0.5 0.5 0.25 0.25\n" if i else ""
|
|
(labels / f"{i}.txt").write_text(content, encoding="utf-8")
|
|
|
|
data_yaml = write_data_yaml(tmp_path / "data.yaml", CLASSES, root=tmp_path)
|
|
|
|
def dataset(split):
|
|
return {
|
|
"dataset": {"name": "film_defect", "img_dir": str(tmp_path / "images" / split),
|
|
"transforms": TRANSFORMS},
|
|
"loader": {"batch_size": 2, "num_workers": 0, "shuffle": split == "train"},
|
|
}
|
|
|
|
config = {
|
|
"classes": str(data_yaml),
|
|
"label_mode": "multilabel",
|
|
"image_size": [32, 32],
|
|
"device": "cpu",
|
|
"seed": 0,
|
|
"model": {
|
|
"name": "defect_net",
|
|
"backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
|
|
"neck": {"name": "upsample", "channels": 8, "num_upsamples": 1},
|
|
"head": {"hidden_channels": 8},
|
|
},
|
|
"data": {"train": dataset("train"), "val": dataset("val")},
|
|
"train": {
|
|
"epochs": 1,
|
|
"amp": False,
|
|
"log_interval": 10,
|
|
"optimizer": {"name": "adamw", "lr": 0.001},
|
|
"swanlab": {"enabled": False},
|
|
},
|
|
}
|
|
config_path = tmp_path / "config.yaml"
|
|
config_path.write_text(yaml.safe_dump(config, allow_unicode=True), encoding="utf-8")
|
|
return config_path, tmp_path
|
|
|
|
|
|
def run(module, argv, monkeypatch):
|
|
monkeypatch.setattr(sys, "argv", argv)
|
|
return module.main()
|
|
|
|
|
|
def test_train_then_eval(project, monkeypatch, tmp_path):
|
|
config_path, root = project
|
|
work_dir = tmp_path / "runs"
|
|
|
|
assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)],
|
|
monkeypatch) == 0
|
|
assert (work_dir / "best.pth").exists()
|
|
assert (work_dir / "config.yaml").exists()
|
|
|
|
metrics_path = tmp_path / "metrics.json"
|
|
assert run(
|
|
eval_mod,
|
|
["eval.py", "--config", str(config_path), "--checkpoint", str(work_dir / "best.pth"),
|
|
"--set", "device=cpu", "--out", str(metrics_path)],
|
|
monkeypatch,
|
|
) == 0
|
|
assert metrics_path.exists()
|
|
|
|
|
|
def test_classes_from_data_yaml_sizes_the_heads(project, monkeypatch, tmp_path):
|
|
"""classes 写成 data.yaml 路径时，num_classes 必须是 2，而不是路径字符串的长度。"""
|
|
config_path, _ = project
|
|
work_dir = tmp_path / "runs2"
|
|
captured = {}
|
|
|
|
real_build = train_mod.build_model
|
|
|
|
def spy(cfg):
|
|
captured["cfg"] = dict(cfg)
|
|
return real_build(cfg)
|
|
|
|
monkeypatch.setattr(train_mod, "build_model", spy)
|
|
|
|
run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)], monkeypatch)
|
|
|
|
assert captured["cfg"]["num_classes"] == len(CLASSES)
|
|
assert captured["cfg"]["label_mode"] == "multilabel"
|
|
|
|
|
|
def test_unavailable_device_fails_loudly(project, monkeypatch, tmp_path):
|
|
"""显式请求机器上没有的后端时必须退出，而不是悄悄退回 CPU。"""
|
|
from utils.device import available_backends
|
|
|
|
missing = next((b for b in ("cuda", "mps") if b not in available_backends()), None)
|
|
if missing is None:
|
|
pytest.skip("这台机器 cuda 与 mps 都可用")
|
|
|
|
assert run(
|
|
train_mod,
|
|
["train.py", "--config", str(config_path_of(project)), "--work-dir", str(tmp_path / "runs3"),
|
|
"--device", missing],
|
|
monkeypatch,
|
|
) == 1
|
|
|
|
|
|
def test_device_flag_overrides_config(project, monkeypatch, tmp_path):
|
|
"""--device 优先于配置里的 device。"""
|
|
config_path, _ = project
|
|
work_dir = tmp_path / "runs4"
|
|
assert run(
|
|
train_mod,
|
|
["train.py", "--config", str(config_path), "--work-dir", str(work_dir), "--device", "cpu"],
|
|
monkeypatch,
|
|
) == 0
|
|
assert (work_dir / "best.pth").exists()
|
|
|
|
|
|
def config_path_of(project):
|
|
return project[0]
|
|
|
|
|
|
def test_resume_takes_hyperparams_from_config(project, monkeypatch, tmp_path):
|
|
"""检测训练续训：lr / weight_decay / 总轮数以 yaml 为准，检查点只提供进度。
|
|
|
|
回归测试：load_checkpoint 直接 optimizer.load_state_dict，把超参连同动量一起还原，
|
|
改了配置再续训会被静默覆盖；而且此前根本没保存调度器状态，学习率调度会从头开始。
|
|
"""
|
|
import math
|
|
|
|
import yaml as yaml_mod
|
|
|
|
config_path, root = project
|
|
cfg = yaml_mod.safe_load(config_path.read_text(encoding="utf-8"))
|
|
cfg["train"].update({"epochs": 4, "optimizer": {"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4},
|
|
"scheduler": {"name": "cosine"}})
|
|
config_path.write_text(yaml_mod.safe_dump(cfg, allow_unicode=True), encoding="utf-8")
|
|
|
|
work_dir = tmp_path / "runs"
|
|
assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)],
|
|
monkeypatch) == 0
|
|
|
|
captured = {}
|
|
real_fit = train_mod.Trainer.fit
|
|
|
|
def probe(self):
|
|
captured["groups"] = [dict(g) for g in self.optimizer.param_groups]
|
|
captured["scheduler"] = self.scheduler
|
|
captured["start_epoch"] = self.start_epoch
|
|
return 0.0
|
|
|
|
monkeypatch.setattr(train_mod.Trainer, "fit", probe)
|
|
assert run(
|
|
train_mod,
|
|
["train.py", "--config", str(config_path), "--work-dir", str(work_dir),
|
|
"--resume", str(work_dir / "last.pth"), "--set",
|
|
"train.epochs=10", "train.optimizer.lr=5e-4", "train.optimizer.weight_decay=0.05"],
|
|
monkeypatch,
|
|
) == 0
|
|
monkeypatch.setattr(train_mod.Trainer, "fit", real_fit)
|
|
|
|
groups = captured["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 分组结构没被冲掉
|
|
|
|
scheduler = captured["scheduler"]
|
|
assert scheduler.T_max == 10
|
|
assert scheduler.last_epoch == captured["start_epoch"] == 4 # 调度进度已快进，不是从 0 开始
|
|
assert groups[0]["lr"] == pytest.approx(5e-4 * (1 + math.cos(math.pi * 4 / 10)) / 2)
|
|
|
|
|
|
def test_checkpoint_carries_scheduler_and_progress(project, monkeypatch, tmp_path):
|
|
"""检查点要带上调度器状态与 global_step，否则续训无从恢复调度进度。"""
|
|
config_path, _ = project
|
|
work_dir = tmp_path / "runs"
|
|
assert run(train_mod, ["train.py", "--config", str(config_path), "--work-dir", str(work_dir)],
|
|
monkeypatch) == 0
|
|
|
|
state = torch.load(work_dir / "last.pth", map_location="cpu", weights_only=False)
|
|
expected = {"model", "optimizer", "scheduler", "epoch", "global_step", "best_metric", "train_cfg"}
|
|
assert expected <= set(state)
|
|
assert state["global_step"] > 0
|
|
assert state["train_cfg"]["optimizer"]["lr"] is not None
|