feat(simclr): 接入 SwanLab 可视化,并补上本地看板缺失的依赖
- SimCLRTrainer 记录 simclr/loss、loss_avg、lr(按 step)与 epoch_loss、 epoch_seconds、gpu_mem_peak_mb(按 epoch),超参为完整 yaml 配置; 正常结束与 KeyboardInterrupt 两条路径都会 finish() - 配置新增 simclr.swanlab(字段同 train.swanlab),默认 mode: local, 不需要登录,日志写到 <work_dir>/swanlog - 依赖改为 swanlab[dashboard]:此前 `swanlab watch` 在干净环境里会报 "Please install the swanboard package",README 里检测训练那节的看板命令 同样受影响,现已随 uv sync 装好 - README 补齐本地看板/云端/离线补传三种用法与端口、局域网参数 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
7d5b6ec423
commit
be6173c746
+36
-2
@@ -18,6 +18,7 @@ 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
|
||||
@@ -45,7 +46,9 @@ class SimCLRTrainer:
|
||||
"""Args:
|
||||
backbone: 输出特征图的卷积骨架(需带 `out_channels`)。
|
||||
dataloader: 每个样本产出 2 个视图的 DataLoader(`UnlabeledImageDataset(views=2)`)。
|
||||
cfg: `epochs / temperature / optimizer / scheduler / amp / log_interval / save_interval` 等。
|
||||
cfg: `epochs / temperature / optimizer / scheduler / amp / log_interval / save_interval`
|
||||
以及 `swanlab`(可视化,字段同 `train.swanlab`)。
|
||||
full_config: 记录到 SwanLab 的超参,一般直接传完整 yaml 配置。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -55,6 +58,7 @@ class SimCLRTrainer:
|
||||
cfg: Mapping[str, Any],
|
||||
device: torch.device,
|
||||
work_dir: str | Path,
|
||||
full_config: dict | None = None,
|
||||
) -> None:
|
||||
self.backbone = backbone
|
||||
self.dataloader = dataloader
|
||||
@@ -81,6 +85,13 @@ class SimCLRTrainer:
|
||||
self.scaler = torch.amp.GradScaler(device.type, enabled=self.amp)
|
||||
self.peak_memory = PeakMemory(device)
|
||||
|
||||
self.global_step = 0
|
||||
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()
|
||||
@@ -101,7 +112,17 @@ class SimCLRTrainer:
|
||||
|
||||
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,
|
||||
@@ -113,11 +134,22 @@ class SimCLRTrainer:
|
||||
|
||||
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,
|
||||
time.time() - start,
|
||||
elapsed,
|
||||
peak,
|
||||
"" if self.peak_memory.exact else "(MPS 采样下界)",
|
||||
)
|
||||
@@ -149,8 +181,10 @@ class SimCLRTrainer:
|
||||
except KeyboardInterrupt:
|
||||
self.save_backbone(epoch)
|
||||
logger.warning("训练被中断,已保存 epoch %d 的权重: %s", epoch, path)
|
||||
self.visualizer.finish()
|
||||
return path
|
||||
|
||||
self.save_backbone(self.epochs - 1)
|
||||
logger.info("自监督骨架已保存: %s", path)
|
||||
self.visualizer.finish()
|
||||
return path
|
||||
|
||||
Reference in New Issue
Block a user