fix(simclr): 周期性落盘,中途中断不再丢失训练成果

此前 fit() 只在全部 epoch 跑完后才保存骨架,一次跑到 epoch 12 被停掉后
13 分钟的训练全部作废。

- 每 save_interval(默认 5)个 epoch 保存一次骨架
- 捕获 KeyboardInterrupt,中断时先落盘再返回
- 聚类脚本改用 resolve_device,自监督分支不再有 GPU 守卫

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-02 21:43:35 +08:00
co-authored by Claude Opus 5
parent 141cb907f2
commit b6f83fcad0
3 changed files with 62 additions and 30 deletions
+13 -13
View File
@@ -4,7 +4,7 @@
# 默认 4 类,ImageNet 预训练特征(不训练,CPU 可跑)
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml
# 改簇数 / 换成在本地数据上自监督训练的特征(需要 Nvidia GPU)
# 改簇数 / 换成在本地数据上自监督训练的特征(Mac 上自动走 MPS)
uv run python scripts/cluster_unlabeled.py --config configs/cluster/film.yaml \
--set num_clusters=6 method=simclr
@@ -40,7 +40,11 @@ from engine.cluster import ( # noqa: E402
)
from models.registry import BACKBONES # noqa: E402
from utils import load_config, seed_everything, setup_logging # noqa: E402
from utils.device import NoNvidiaGpuError, require_cuda, resolve_device # noqa: E402
from utils.device import ( # noqa: E402
DeviceNotAvailableError,
describe_device,
resolve_device,
)
logger = None
@@ -49,7 +53,7 @@ def build_encoder(cfg, device, work_dir):
"""按 `method` 拿一个特征提取器。
pretrained: ImageNet 预训练卷积骨架,纯推理,CPU 也能跑。用顶层的 `backbone`。
simclr: 在本目录数据上自监督训练一个项目自己的卷积骨架——属于训练,先过 GPU 守卫。
simclr: 在本目录数据上自监督训练一个项目自己的卷积骨架。
用 `simclr.backbone`,与 pretrained 的骨架配置各自独立:两者的字段完全不同
(arch/pretrained vs widths/depths),共用一个块会在 `_base_` 合并时串味。
"""
@@ -62,7 +66,6 @@ def build_encoder(cfg, device, work_dir):
if method != "simclr":
raise ValueError(f"未知 method: {method},可选 pretrained / simclr")
require_cuda() # 自监督预训练是训练,禁止在无 Nvidia GPU 的设备上跑
from torch.utils.data import DataLoader
from engine.simclr import SimCLRTrainer
@@ -200,6 +203,7 @@ def main() -> int:
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(
@@ -220,26 +224,22 @@ def main() -> int:
num_clusters = int(cfg.get("num_clusters", 4))
try:
device = resolve_device(str(cfg.get("device", "cuda")))
except NoNvidiaGpuError as exc:
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,device=%s",
"扫描到 %d 张图,来自 %d 个目录;簇数 %d,method=%s,设备 %s",
len(dataset),
len({dataset.source_dir(path) for path in dataset.images}),
num_clusters,
cfg.get("method", "pretrained"),
device,
describe_device(device),
)
try:
encoder = build_encoder(cfg, device, work_dir)
except NoNvidiaGpuError as exc:
logger.error("%s", exc)
return 1
encoder = build_encoder(cfg, device, work_dir)
from torch.utils.data import DataLoader