Shaun ZhouandClaude Opus 5
ef64e5bf77
- trainer: AMP、梯度裁剪、cosine/step 调度、last/best 检查点，按 monitor 指标挑选 best - losses: penalty-reduced focal loss（heatmap）+ 掩码 L1（wh/offset）+ 分类损失 （多类 BCE / 单类 CE，随 label_mode 切换） - metrics: 定位 IoU 匹配下的 P/R/F1 与 mean IoU；分类多标签 micro-F1 与单标签 accuracy， 两种模式都输出 cls_f1，故 train.monitor 无需跟着改 - optim: BatchNorm 与 bias 不做 weight decay - visualizer: SwanLab 记录损失/学习率/验证指标/显存峰值，未安装或离线时降级为警告 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
47 lines
1.1 KiB
Python
47 lines
1.1 KiB
Python
"""检查点存取。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
def save_checkpoint(
|
|
path: str | Path,
|
|
model: nn.Module,
|
|
optimizer: torch.optim.Optimizer | None = None,
|
|
epoch: int = 0,
|
|
best_metric: float | None = None,
|
|
config: dict | None = None,
|
|
) -> Path:
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
torch.save(
|
|
{
|
|
"model": model.state_dict(),
|
|
"optimizer": optimizer.state_dict() if optimizer else None,
|
|
"epoch": epoch,
|
|
"best_metric": best_metric,
|
|
"config": config,
|
|
},
|
|
path,
|
|
)
|
|
return path
|
|
|
|
|
|
def load_checkpoint(
|
|
path: str | Path,
|
|
model: nn.Module,
|
|
optimizer: torch.optim.Optimizer | None = None,
|
|
map_location: Any = "cpu",
|
|
strict: bool = True,
|
|
) -> dict:
|
|
state = torch.load(Path(path), map_location=map_location, weights_only=False)
|
|
model.load_state_dict(state["model"], strict=strict)
|
|
if optimizer is not None and state.get("optimizer"):
|
|
optimizer.load_state_dict(state["optimizer"])
|
|
return state
|