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>
88 lines
3.2 KiB
Python
88 lines
3.2 KiB
Python
"""评估：跑一遍验证集，输出分类与定位指标。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import torch
|
|
from torch import nn
|
|
from torch.utils.data import DataLoader
|
|
|
|
from engine.metrics import AverageMeter, DetectionMetric, MultiLabelMetric, SingleLabelMetric
|
|
from models.postprocess import decode_heatmap
|
|
from utils.labels import MULTILABEL
|
|
from utils.logging import get_logger
|
|
|
|
logger = get_logger("engine.evaluator")
|
|
|
|
|
|
class Evaluator:
|
|
def __init__(
|
|
self,
|
|
model: nn.Module,
|
|
dataloader: DataLoader,
|
|
device: torch.device,
|
|
criterion: nn.Module | None = None,
|
|
score_threshold: float = 0.3,
|
|
iou_threshold: float = 0.5,
|
|
topk: int = 50,
|
|
) -> None:
|
|
self.model = model
|
|
self.dataloader = dataloader
|
|
self.device = device
|
|
self.criterion = criterion
|
|
self.score_threshold = score_threshold
|
|
self.topk = topk
|
|
self.det_metric = DetectionMetric(iou_threshold)
|
|
# 分类指标按模型的 label_mode 选择：多类用多标签指标，单类用互斥分类指标。
|
|
self.label_mode = getattr(model, "label_mode", MULTILABEL)
|
|
if self.label_mode == MULTILABEL:
|
|
self.cls_metric: MultiLabelMetric | SingleLabelMetric = MultiLabelMetric()
|
|
else:
|
|
num_outputs = model.cls_head.num_outputs if model.cls_head is not None else 1
|
|
self.cls_metric = SingleLabelMetric(num_outputs)
|
|
|
|
@torch.no_grad()
|
|
def evaluate(self) -> dict[str, float]:
|
|
self.model.eval()
|
|
self.det_metric.reset()
|
|
self.cls_metric.reset()
|
|
loss_meter = AverageMeter()
|
|
|
|
stride = getattr(self.model, "out_stride", 4)
|
|
for batch in self.dataloader:
|
|
images = batch["images"].to(self.device, non_blocking=True)
|
|
outputs = self.model(images)
|
|
|
|
if self.criterion is not None:
|
|
batch_on_device = _to_device(batch, self.device)
|
|
losses = self.criterion(outputs, batch_on_device)
|
|
loss_meter.update(float(losses["total"]), images.shape[0])
|
|
|
|
detections = decode_heatmap(
|
|
outputs["heatmap"],
|
|
outputs["wh"],
|
|
outputs["offset"],
|
|
stride=stride,
|
|
topk=self.topk,
|
|
score_threshold=self.score_threshold,
|
|
)
|
|
for i, det in enumerate(detections):
|
|
self.det_metric.update(det["boxes"], det["labels"], batch["boxes"][i], batch["labels"][i])
|
|
|
|
if "logits" in outputs:
|
|
self.cls_metric.update(outputs["logits"], batch["image_labels"])
|
|
|
|
metrics: dict[str, float] = {**self.det_metric.compute(), **self.cls_metric.compute()}
|
|
if self.criterion is not None:
|
|
metrics["loss"] = loss_meter.avg
|
|
return metrics
|
|
|
|
|
|
def _to_device(batch: dict, device: torch.device) -> dict:
|
|
moved: dict[str, Any] = dict(batch)
|
|
moved["images"] = batch["images"].to(device, non_blocking=True)
|
|
moved["target"] = {k: v.to(device, non_blocking=True) for k, v in batch["target"].items()}
|
|
moved["image_labels"] = batch["image_labels"].to(device, non_blocking=True)
|
|
return moved
|