xzhou/ModelTrainEval
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
此前 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
Download Patch File Download Diff File Expand all files Collapse all files
configs/cluster/film.yaml
+5 -4
View File
@@ -22,12 +22,12 @@
work_dir: runs/cluster
seed: 42
device: cuda # method=pretrained 时可以 --set device=cpu
device: auto # auto 按 cuda > mps > cpu 挑最快的可用后端；也可写 cuda / mps / cpu
num_clusters: 6 # 想分几类改这里（或 --num-clusters 6）
num_clusters: 4 # 想分几类改这里（或 --num-clusters 6）
# pretrained: ImageNet 预训练卷积骨架直接抽特征，不训练，CPU 也能跑，几分钟出结果。
# simclr: 在本目录数据上自监督训练项目自己的卷积骨架，域适配更好，需要 Nvidia GPU。
# simclr: 在本目录数据上自监督训练项目自己的卷积骨架，域适配更好；Mac 上走 MPS 即可。
method: pretrained
# method=pretrained 用的骨架（method=simclr 的骨架配在 simclr.backbone，两者字段完全不同，
@@ -97,7 +97,8 @@ simclr:
epochs: 100
batch_size: 256
temperature: 0.2 # 越小越关注难负样本
amp: true
amp: auto # CUDA 上开、MPS 上关
save_interval: 5 # 每 5 个 epoch 落一次盘，中途中断也能拿到权重
log_interval: 20
checkpoint: null # 填已有的 simclr_backbone.pth 可跳过重复训练
projection: {hidden_channels: 256, out_channels: 128}
engine/simclr.py
+44 -13
View File
@@ -1,7 +1,7 @@
"""SimCLR 自监督预训练：在无标注的薄膜图上学一套特征，供聚类使用。
骨架仍然是项目里的纯卷积 backbone（`tiny_conv` 等），预训练完可以直接拿去
初始化检测模型。属于"训练"，因此调用方必须先过 `utils.device.require_cuda()`。
初始化检测模型。支持 cuda / mps / cpu，用 `device: auto` 自动挑最快的后端。
"""
from __future__ import annotations
@@ -19,6 +19,7 @@ from torch.utils.data import DataLoader
from engine.metrics import AverageMeter
from engine.optim import build_optimizer, build_scheduler
from models.heads.projection import ProjectionHead
from utils.device import PeakMemory, resolve_amp, synchronize
from utils.logging import get_logger
logger = get_logger("engine.simclr")
@@ -44,7 +45,7 @@ class SimCLRTrainer:
"""Args:
backbone: 输出特征图的卷积骨架（需带 `out_channels`）。
dataloader: 每个样本产出 2 个视图的 DataLoader（`UnlabeledImageDataset(views=2)`）。
cfg: `epochs / temperature / optimizer / scheduler / amp / log_interval` 等。
cfg: `epochs / temperature / optimizer / scheduler / amp / log_interval / save_interval` 等。
"""
def __init__(
@@ -65,6 +66,8 @@ class SimCLRTrainer:
self.epochs = int(self.cfg.get("epochs", 50))
self.temperature = float(self.cfg.get("temperature", 0.2))
self.log_interval = int(self.cfg.get("log_interval", 20))
# 每隔若干 epoch 落一次盘：自监督动辄跑几十分钟，中途被打断不该把训练成果全丢掉
self.save_interval = int(self.cfg.get("save_interval", 5))
self.projector = ProjectionHead(
backbone.out_channels, **dict(self.cfg.get("projection", {}))
@@ -74,8 +77,9 @@ class SimCLRTrainer:
self.model, self.cfg.get("optimizer", {"name": "adamw", "lr": 1e-3, "weight_decay": 1e-4})
)
self.scheduler = build_scheduler(self.optimizer, self.cfg.get("scheduler"), self.epochs)
self.amp = bool(self.cfg.get("amp", True)) and device.type == "cuda"
self.scaler = torch.amp.GradScaler("cuda", enabled=self.amp)
self.amp = resolve_amp(self.cfg.get("amp", "auto"), device)
self.scaler = torch.amp.GradScaler(device.type, enabled=self.amp)
self.peak_memory = PeakMemory(device)
def train_one_epoch(self, epoch: int) -> float:
self.model.train()
@@ -85,7 +89,7 @@ class SimCLRTrainer:
for step, batch in enumerate(self.dataloader):
view1, view2 = (view.to(self.device, non_blocking=True) for view in batch["views"][:2])
with torch.amp.autocast("cuda", enabled=self.amp):
with torch.amp.autocast(self.device.type, enabled=self.amp):
z1 = self.projector(self.backbone(view1))
z2 = self.projector(self.backbone(view2))
loss = nt_xent_loss(z1, z2, self.temperature)
@@ -95,6 +99,7 @@ class SimCLRTrainer:
self.scaler.step(self.optimizer)
self.scaler.update()
self.peak_memory.update()
meter.update(float(loss.detach()), view1.shape[0])
if step % self.log_interval == 0:
logger.info(
@@ -106,20 +111,46 @@ class SimCLRTrainer:
self.optimizer.param_groups[0]["lr"],
)
logger.info("simclr epoch %d 完成 loss=%.4f 耗时 %.1fs", epoch, meter.avg, time.time() - start)
synchronize(self.device)
peak = self.peak_memory.peak_mb
logger.info(
"simclr epoch %d 完成 loss=%.4f 耗时 %.1fs 显存峰值 %.0f MB%s",
epoch,
meter.avg,
time.time() - start,
peak,
"" if self.peak_memory.exact else "（MPS 采样下界）",
)
return meter.avg
def save_backbone(self, epoch: int) -> Path:
"""只存骨架权重，不含投影头——投影头训练成对增强不变，会丢掉下游有用的信息。"""
path = self.work_dir / "simclr_backbone.pth"
torch.save({"model": self.backbone.state_dict(), "epochs": epoch + 1}, path)
return path
def fit(self) -> Path:
"""训练并保存骨架权重（不含投影头），返回权重路径。"""
"""训练并保存骨架权重（不含投影头），返回权重路径。
每 `save_interval` 个 epoch 落一次盘，中途被 Ctrl-C 或 kill 掉也能拿到当前权重继续用。
"""
if len(self.dataloader) < 2:
raise ValueError("SimCLR 需要多个 batch，请调小 batch_size 或增加样本")
for epoch in range(self.epochs):
self.train_one_epoch(epoch)
if self.scheduler is not None:
self.scheduler.step()
path = self.work_dir / "simclr_backbone.pth"
torch.save({"model": self.backbone.state_dict(), "epochs": self.epochs}, path)
try:
for epoch in range(self.epochs):
self.train_one_epoch(epoch)
if self.scheduler is not None:
self.scheduler.step()
if self.save_interval and (epoch + 1) % self.save_interval == 0:
self.save_backbone(epoch)
logger.info("已保存中间权重（epoch %d）: %s", epoch, path)
except KeyboardInterrupt:
self.save_backbone(epoch)
logger.warning("训练被中断，已保存 epoch %d 的权重: %s", epoch, path)
return path
self.save_backbone(self.epochs - 1)
logger.info("自监督骨架已保存: %s", path)
return path
scripts/cluster_unlabeled.py
+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