Shaun ZhouandClaude Opus 5
345dee50da
此前 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>
275 lines
11 KiB
Python
275 lines
11 KiB
Python
"""SimCLR 自监督预训练：在无标注的薄膜图上学一套特征，供聚类使用。
|
||
|
||
骨架仍然是项目里的纯卷积 backbone（`tiny_conv` 等），预训练完可以直接拿去
|
||
初始化检测模型。支持 cuda / mps / cpu，用 `device: auto` 自动挑最快的后端。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
from collections.abc import Mapping
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import torch
|
||
from torch import nn
|
||
from torch.nn import functional as F
|
||
from torch.utils.data import DataLoader
|
||
|
||
from engine.metrics import AverageMeter
|
||
from engine.optim import build_optimizer, build_scheduler
|
||
from engine.visualizer import SwanLabLogger
|
||
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")
|
||
|
||
|
||
def nt_xent_loss(z1: torch.Tensor, z2: torch.Tensor, temperature: float = 0.2) -> torch.Tensor:
|
||
"""NT-Xent：同一张图的两个视图互为正样本，batch 内其余 2N-2 个都是负样本。"""
|
||
batch_size = z1.shape[0]
|
||
z = F.normalize(torch.cat([z1, z2]), dim=1)
|
||
similarity = z @ z.t() / temperature
|
||
similarity.fill_diagonal_(float("-inf")) # 排除自己和自己
|
||
|
||
targets = torch.cat(
|
||
[
|
||
torch.arange(batch_size, 2 * batch_size, device=z.device),
|
||
torch.arange(0, batch_size, device=z.device),
|
||
]
|
||
)
|
||
return F.cross_entropy(similarity, targets)
|
||
|
||
|
||
class SimCLRTrainer:
|
||
"""Args:
|
||
backbone: 输出特征图的卷积骨架（需带 `out_channels`）。
|
||
dataloader: 每个样本产出 2 个视图的 DataLoader（`UnlabeledImageDataset(views=2)`）。
|
||
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__(
|
||
self,
|
||
backbone: nn.Module,
|
||
dataloader: DataLoader,
|
||
cfg: Mapping[str, Any],
|
||
device: torch.device,
|
||
work_dir: str | Path,
|
||
full_config: dict | None = None,
|
||
resume: str | Path | None = None,
|
||
) -> None:
|
||
self.backbone = backbone
|
||
self.dataloader = dataloader
|
||
self.cfg = dict(cfg)
|
||
self.device = device
|
||
self.work_dir = Path(work_dir)
|
||
self.work_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
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", {}))
|
||
).to(device)
|
||
self.model = nn.ModuleDict({"backbone": backbone, "projector": self.projector})
|
||
self.optimizer = build_optimizer(
|
||
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 = resolve_amp(self.cfg.get("amp", "auto"), device)
|
||
self.scaler = torch.amp.GradScaler(device.type, enabled=self.amp)
|
||
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,
|
||
**dict(self.cfg.get("swanlab", {})),
|
||
)
|
||
|
||
def train_one_epoch(self, epoch: int) -> float:
|
||
self.model.train()
|
||
meter = AverageMeter()
|
||
start = time.time()
|
||
|
||
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(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)
|
||
|
||
self.optimizer.zero_grad(set_to_none=True)
|
||
self.scaler.scale(loss).backward()
|
||
self.scaler.step(self.optimizer)
|
||
self.scaler.update()
|
||
|
||
self.peak_memory.update()
|
||
meter.update(float(loss.detach()), view1.shape[0])
|
||
self.global_step += 1
|
||
if step % self.log_interval == 0:
|
||
self.visualizer.log(
|
||
{
|
||
"loss": float(loss.detach()),
|
||
"loss_avg": meter.avg,
|
||
"lr": self.optimizer.param_groups[0]["lr"],
|
||
},
|
||
step=self.global_step,
|
||
prefix="simclr/",
|
||
)
|
||
logger.info(
|
||
"simclr epoch %d [%d/%d] loss=%.4f lr=%.2e",
|
||
epoch,
|
||
step,
|
||
len(self.dataloader),
|
||
meter.avg,
|
||
self.optimizer.param_groups[0]["lr"],
|
||
)
|
||
|
||
synchronize(self.device)
|
||
peak = self.peak_memory.peak_mb
|
||
elapsed = time.time() - start
|
||
self.visualizer.log(
|
||
{
|
||
"epoch": float(epoch),
|
||
"epoch_loss": meter.avg,
|
||
"epoch_seconds": elapsed,
|
||
"gpu_mem_peak_mb": peak,
|
||
},
|
||
step=self.global_step,
|
||
prefix="simclr/",
|
||
)
|
||
logger.info(
|
||
"simclr epoch %d 完成 loss=%.4f 耗时 %.1fs 显存峰值 %.0f MB%s",
|
||
epoch,
|
||
meter.avg,
|
||
elapsed,
|
||
peak,
|
||
"" if self.peak_memory.exact else "（MPS 采样下界）",
|
||
)
|
||
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 save_state(self, epoch: int) -> Path:
|
||
"""完整训练状态，供 `resume` 断点续训。
|
||
|
||
比 `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.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._checkpoint(epoch)
|
||
logger.info("已保存中间权重（epoch %d）: %s", epoch, path)
|
||
except KeyboardInterrupt:
|
||
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._checkpoint(self.epochs - 1)
|
||
logger.info("自监督骨架已保存: %s", path)
|
||
self.visualizer.finish()
|
||
return path
|