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>
62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
"""优化器与学习率调度。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import Mapping
|
||
from typing import Any
|
||
|
||
import torch
|
||
from torch import nn
|
||
|
||
|
||
def build_optimizer(model: nn.Module, cfg: Mapping[str, Any]) -> torch.optim.Optimizer:
|
||
"""`cfg` 形如 `{name: sgd, lr: 0.01, momentum: 0.9, weight_decay: 1e-4}`。
|
||
|
||
BatchNorm 与 bias 不做 weight decay——小数据集上这项影响明显。
|
||
"""
|
||
cfg = dict(cfg)
|
||
name = str(cfg.pop("name", "adamw")).lower()
|
||
weight_decay = float(cfg.pop("weight_decay", 0.0))
|
||
|
||
decay, no_decay = [], []
|
||
for module in model.modules():
|
||
for param_name, param in module.named_parameters(recurse=False):
|
||
if not param.requires_grad:
|
||
continue
|
||
if param_name == "bias" or isinstance(module, nn.modules.batchnorm._BatchNorm):
|
||
no_decay.append(param)
|
||
else:
|
||
decay.append(param)
|
||
groups = [
|
||
{"params": decay, "weight_decay": weight_decay},
|
||
{"params": no_decay, "weight_decay": 0.0},
|
||
]
|
||
|
||
if name == "sgd":
|
||
cfg.setdefault("momentum", 0.9)
|
||
cfg.setdefault("nesterov", True)
|
||
return torch.optim.SGD(groups, **cfg)
|
||
if name == "adam":
|
||
return torch.optim.Adam(groups, **cfg)
|
||
if name == "adamw":
|
||
return torch.optim.AdamW(groups, **cfg)
|
||
raise ValueError(f"未知优化器: {name}")
|
||
|
||
|
||
def build_scheduler(optimizer: torch.optim.Optimizer, cfg: Mapping[str, Any] | None, epochs: int):
|
||
"""返回按 epoch 步进的 scheduler；`cfg` 为空则不使用。"""
|
||
if not cfg:
|
||
return None
|
||
cfg = dict(cfg)
|
||
name = str(cfg.pop("name", "cosine")).lower()
|
||
if name == "cosine":
|
||
cfg.setdefault("T_max", epochs)
|
||
return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **cfg)
|
||
if name == "step":
|
||
cfg.setdefault("step_size", max(epochs // 3, 1))
|
||
return torch.optim.lr_scheduler.StepLR(optimizer, **cfg)
|
||
if name == "multistep":
|
||
cfg.setdefault("milestones", [int(epochs * 0.6), int(epochs * 0.85)])
|
||
return torch.optim.lr_scheduler.MultiStepLR(optimizer, **cfg)
|
||
raise ValueError(f"未知调度器: {name}")
|