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:
co-authored by
Claude Opus 5
parent
141cb907f2
commit
b6f83fcad0
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user