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>
100 lines
3.5 KiB
Python
Executable File
100 lines
3.5 KiB
Python
Executable File
#!/usr/bin/env python
|
||
"""训练入口。
|
||
|
||
用法:
|
||
python train.py --config configs/train/baseline.yaml
|
||
python train.py --config configs/train/baseline.yaml --set train.epochs=10 train.optimizer.lr=5e-4
|
||
python train.py --config configs/train/baseline.yaml --device mps
|
||
|
||
设备由配置的 `device` 决定，默认 `auto`：按 cuda > mps > cpu 挑最快的可用后端。
|
||
显式指定但不可用时直接报错，不会悄悄退回 CPU。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
from datasets import build_dataloader, build_dataset, load_classes
|
||
from engine import Trainer
|
||
from models import build_model
|
||
from utils import check_label_mode, load_config, seed_everything, setup_logging
|
||
from utils.device import DeviceNotAvailableError, describe_device, resolve_device
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(description="薄膜缺陷分类定位模型训练")
|
||
parser.add_argument("--config", required=True, help="yaml 配置路径")
|
||
parser.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE", help="覆盖配置项")
|
||
parser.add_argument("--work-dir", default=None, help="覆盖配置中的 work_dir")
|
||
parser.add_argument("--device", default=None, help="覆盖配置中的 device：auto / cuda / mps / cpu")
|
||
parser.add_argument("--resume", default=None, help="从检查点恢复训练")
|
||
return parser.parse_args()
|
||
|
||
|
||
def main() -> int:
|
||
args = parse_args()
|
||
cfg = load_config(args.config, args.set)
|
||
|
||
work_dir = Path(args.work_dir or cfg.get("work_dir", "runs/default"))
|
||
work_dir.mkdir(parents=True, exist_ok=True)
|
||
logger = setup_logging(work_dir / "train.log")
|
||
|
||
try:
|
||
device = resolve_device(args.device or cfg.get("device", "auto"))
|
||
except DeviceNotAvailableError as exc:
|
||
logger.error("%s", exc)
|
||
return 1
|
||
logger.info("训练设备: %s", describe_device(device))
|
||
|
||
seed_everything(int(cfg.get("seed", 42)), bool(cfg.get("deterministic", False)))
|
||
|
||
classes = load_classes(cfg["classes"])
|
||
# label_mode 由 yaml 决定，模型与数据集必须一致，这里统一下发。
|
||
label_mode = check_label_mode(str(cfg.get("label_mode", "multilabel")))
|
||
model_cfg = dict(cfg["model"])
|
||
model_cfg.setdefault("num_classes", len(classes))
|
||
model_cfg.setdefault("label_mode", label_mode)
|
||
model = build_model(model_cfg).to(device)
|
||
logger.info(
|
||
"模型 %s 输出 stride=%d，参数量 %.2fM",
|
||
model_cfg["name"],
|
||
model.out_stride,
|
||
sum(p.numel() for p in model.parameters()) / 1e6,
|
||
)
|
||
|
||
loaders = {}
|
||
for split in ("train", "val"):
|
||
split_cfg = cfg["data"].get(split)
|
||
if not split_cfg:
|
||
loaders[split] = None
|
||
continue
|
||
dataset = build_dataset(
|
||
split_cfg["dataset"], classes=classes, stride=model.out_stride, label_mode=label_mode
|
||
)
|
||
loaders[split] = build_dataloader(dataset, **dict(split_cfg.get("loader", {})))
|
||
logger.info("%s 集样本数: %d", split, len(dataset))
|
||
|
||
trainer = Trainer(
|
||
model=model,
|
||
train_loader=loaders["train"],
|
||
val_loader=loaders["val"],
|
||
cfg=cfg.get("train", {}),
|
||
device=device,
|
||
work_dir=work_dir,
|
||
full_config=cfg.to_dict(),
|
||
)
|
||
|
||
if args.resume:
|
||
trainer.resume(args.resume)
|
||
|
||
cfg.dump(work_dir / "config.yaml")
|
||
best = trainer.fit()
|
||
logger.info("训练结束，最佳 %s=%.4f", trainer.monitor, best)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|