feat(simclr): 支持断点续训，恢复优化器与调度器状态
此前 simclr.checkpoint 的语义是"有权重就跳过训练"，没有真正的续训： 中途停掉再启动只能从头再来，或者放弃已训的轮次。 - 新增 simclr_last.pth：骨架 + 投影头 + 优化器 + 调度器 + scaler + epoch + step， 与给下游用的 simclr_backbone.pth 并存，两者在同一时机落盘 - 新增 simclr.resume：载入完整状态从下一个 epoch 接着练。只恢复权重不恢复动量与 cosine 进度的话，续训时损失会先反弹一段 - resume 与 checkpoint 互斥（同时给直接报错）；拿只含骨架的 simclr_backbone.pth 去 resume 会给出可操作的报错，提示改用 checkpoint - 续训轮数已达计划值时不再训练，只打警告 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
be6173c746
commit
345dee50da
+93
-9
@@ -49,6 +49,9 @@ class SimCLRTrainer:
|
||||
cfg: `epochs / temperature / optimizer / scheduler / amp / log_interval / save_interval`
|
||||
以及 `swanlab`（可视化，字段同 `train.swanlab`）。
|
||||
full_config: 记录到 SwanLab 的超参，一般直接传完整 yaml 配置。
|
||||
resume: 训练状态文件（`simclr_last.pth`）路径，从上次中断处接着练。
|
||||
与 `simclr.checkpoint` 不同：checkpoint 是"有权重就跳过训练直接聚类"，
|
||||
resume 是"载入骨架 + 投影头 + 优化器 + 调度器 + 起始 epoch 继续训练"。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -59,6 +62,7 @@ class SimCLRTrainer:
|
||||
device: torch.device,
|
||||
work_dir: str | Path,
|
||||
full_config: dict | None = None,
|
||||
resume: str | Path | None = None,
|
||||
) -> None:
|
||||
self.backbone = backbone
|
||||
self.dataloader = dataloader
|
||||
@@ -86,6 +90,10 @@ class SimCLRTrainer:
|
||||
self.peak_memory = PeakMemory(device)
|
||||
|
||||
self.global_step = 0
|
||||
self.start_epoch = 0
|
||||
if resume:
|
||||
self.load_state(resume)
|
||||
|
||||
self.visualizer = SwanLabLogger(
|
||||
work_dir=self.work_dir,
|
||||
config=full_config,
|
||||
@@ -156,35 +164,111 @@ class SimCLRTrainer:
|
||||
return meter.avg
|
||||
|
||||
def save_backbone(self, epoch: int) -> Path:
|
||||
"""只存骨架权重，不含投影头——投影头训练成对增强不变，会丢掉下游有用的信息。"""
|
||||
"""只存骨架权重，不含投影头——投影头训练成对增强不变，会丢掉下游有用的信息。
|
||||
|
||||
这是给下游用的产物（聚类、初始化 DefectNet）；要接着训练请用 `simclr_last.pth`。
|
||||
"""
|
||||
path = self.work_dir / "simclr_backbone.pth"
|
||||
torch.save({"model": self.backbone.state_dict(), "epochs": epoch + 1}, path)
|
||||
return path
|
||||
|
||||
def fit(self) -> Path:
|
||||
"""训练并保存骨架权重（不含投影头），返回权重路径。
|
||||
def save_state(self, epoch: int) -> Path:
|
||||
"""完整训练状态，供 `resume` 断点续训。
|
||||
|
||||
每 `save_interval` 个 epoch 落一次盘，中途被 Ctrl-C 或 kill 掉也能拿到当前权重继续用。
|
||||
比 `simclr_backbone.pth` 多了投影头、优化器、调度器与 scaler——少了这些，
|
||||
续训等于把动量和学习率进度清零，损失会先反弹一段。
|
||||
"""
|
||||
path = self.work_dir / "simclr_last.pth"
|
||||
torch.save(
|
||||
{
|
||||
"backbone": self.backbone.state_dict(),
|
||||
"projector": self.projector.state_dict(),
|
||||
"optimizer": self.optimizer.state_dict(),
|
||||
"scheduler": self.scheduler.state_dict() if self.scheduler else None,
|
||||
"scaler": self.scaler.state_dict(),
|
||||
"epoch": epoch,
|
||||
"global_step": self.global_step,
|
||||
"epochs_planned": self.epochs,
|
||||
},
|
||||
path,
|
||||
)
|
||||
return path
|
||||
|
||||
def load_state(self, path: str | Path) -> None:
|
||||
"""载入 `save_state` 写出的完整状态，从下一个 epoch 接着练。"""
|
||||
path = Path(path)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"续训状态文件不存在: {path}")
|
||||
|
||||
state = torch.load(path, map_location=self.device, weights_only=False)
|
||||
if "projector" not in state:
|
||||
raise ValueError(
|
||||
f"{path} 只含骨架权重，无法续训。续训要用 simclr_last.pth；"
|
||||
"若只是想跳过训练直接聚类，请改用 simclr.checkpoint。"
|
||||
)
|
||||
|
||||
self.backbone.load_state_dict(state["backbone"])
|
||||
self.projector.load_state_dict(state["projector"])
|
||||
self.optimizer.load_state_dict(state["optimizer"])
|
||||
if self.scheduler is not None and state.get("scheduler"):
|
||||
self.scheduler.load_state_dict(state["scheduler"])
|
||||
if state.get("scaler"):
|
||||
self.scaler.load_state_dict(state["scaler"])
|
||||
|
||||
self.start_epoch = int(state["epoch"]) + 1
|
||||
self.global_step = int(state.get("global_step", 0))
|
||||
logger.info(
|
||||
"已载入训练状态 %s：从 epoch %d 接着练到 %d", path, self.start_epoch, self.epochs
|
||||
)
|
||||
if self.start_epoch >= self.epochs:
|
||||
logger.warning(
|
||||
"状态里已训到 epoch %d，不小于计划的 %d 轮——不会再训练。"
|
||||
"要继续请调大 simclr.epochs。",
|
||||
self.start_epoch,
|
||||
self.epochs,
|
||||
)
|
||||
|
||||
def _checkpoint(self, epoch: int) -> None:
|
||||
"""同时落两份：给下游用的骨架，和给续训用的完整状态。"""
|
||||
self.save_backbone(epoch)
|
||||
self.save_state(epoch)
|
||||
|
||||
def fit(self) -> Path:
|
||||
"""训练并保存骨架权重（不含投影头），返回骨架路径。
|
||||
|
||||
每 `save_interval` 个 epoch 落一次盘（Ctrl-C 也会先落盘），
|
||||
同时写 `simclr_backbone.pth`（下游用）与 `simclr_last.pth`（续训用）。
|
||||
"""
|
||||
if len(self.dataloader) < 2:
|
||||
raise ValueError("SimCLR 需要多个 batch，请调小 batch_size 或增加样本")
|
||||
|
||||
path = self.work_dir / "simclr_backbone.pth"
|
||||
if self.start_epoch >= self.epochs:
|
||||
self.save_backbone(self.epochs - 1)
|
||||
self.visualizer.finish()
|
||||
return path
|
||||
|
||||
epoch = self.start_epoch
|
||||
try:
|
||||
for epoch in range(self.epochs):
|
||||
for epoch in range(self.start_epoch, 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)
|
||||
self._checkpoint(epoch)
|
||||
logger.info("已保存中间权重（epoch %d）: %s", epoch, path)
|
||||
except KeyboardInterrupt:
|
||||
self.save_backbone(epoch)
|
||||
logger.warning("训练被中断，已保存 epoch %d 的权重: %s", epoch, path)
|
||||
self._checkpoint(epoch)
|
||||
logger.warning(
|
||||
"训练被中断，已保存 epoch %d 的权重: %s；续训用 --set simclr.resume=%s",
|
||||
epoch,
|
||||
path,
|
||||
self.work_dir / "simclr_last.pth",
|
||||
)
|
||||
self.visualizer.finish()
|
||||
return path
|
||||
|
||||
self.save_backbone(self.epochs - 1)
|
||||
self._checkpoint(self.epochs - 1)
|
||||
logger.info("自监督骨架已保存: %s", path)
|
||||
self.visualizer.finish()
|
||||
return path
|
||||
|
||||
Reference in New Issue
Block a user