引言——YOLO为何改变了一切
在 YOLO 出现之前，计算机无法像人类那样“看待”世界。
目标检测系统曾经谨慎、缓慢且分散。它们最初提出的区域是 可能 首先检测包含物体的区域，然后分别对每个区域进行分类。检测结果有效——但感觉就像一块一块地拼拼图。
2015年，YOLO——你只看一次—提出了一个激进的想法：
如果我们在一次前向传播中就能检测到所有问题呢？
YOLO 并没有采用多阶段检测，而是将其视为一个整体。 单一回归问题 从像素到边界框和类别概率。
本指南将逐步介绍如何操作 使用 PyTorch 从零开始完全实现 YOLO，涵盖：
数学公式
网络架构
目标编码
损失实施
在 COCO 风格数据上进行训练
mAP评估
可视化与调试
使用 NMS 进行推理
锚盒扩展
1）YOLO 的含义（以及我们将要构建的东西）
YOLO（你只看一次） 是一系列目标检测模型，用于预测边界框和类别概率。 一次前传与传统的多阶段流程（候选集→优化→分类）不同，YOLO风格的检测器是 密集预测器：它们在多个位置和尺度上预测候选框，然后对其进行筛选。
YOLO 型探测器的发展经历了两个“时代”：
YOLOv1 风格（网格单元格，无锚点）每个网格单元直接预测几个方框。
基于锚点的 YOLO（YOLOv2/3 及许多衍生版本）：每个网格单元预测相对于预定义锚形状的偏移量；多个尺度预测小/中/大物体。
我们将实施
A 现代的、基于锚点的YOLO式探测器 使用：
多鳞片头部（例如，3个鳞片）
锚点匹配（目标分配）
损失函数采用盒回归、目标性检验和分类
解码 + NMS
mAP评估
COCO/自定义数据集训练支持
我们将保持架构的易懂性，而不是过于复杂。以后您可以轻松地升级到更大的主干网。
2) 边界框格式和坐标系
你必须保持一致。大多数训练错误都源于方框格式混乱。
常见包装盒格式：
XYXY:
(x1, y1, x2, y2)左上角和右下角XYWH:
(cx, cy, w, h)中心和尺寸归一化：坐标
[0, 1]相对于图像尺寸绝对像素坐标
推荐的内部惯例
将数据集注释存储为 绝对XYXY 以像素为单位。
仅在需要时才转换为标准化格式，但要保留一个标准格式。
为什么XYXY很棒：
交集/并集运算很简单。
将图像限制在边界内很简单。
3) IoU、GIoU、DIoU、CIoU
IoU（交并比） 是标准重叠度指标：
IoU=∣A∩B∣/∣A∪B∣
但 IoU 存在一个问题：如果检测框不重叠，则 IoU = 0，梯度可能很弱。现代检测器通常使用改进的回归损失：
GIoU：根据最小外接框，对不重叠的框增加惩罚
DIoU：惩罚中心距离
CIoU：DIoU + 宽高比一致性
实用规则:
如果您想要一个强默认值： CIoU 用于盒回归。
如果你想要更简单的方案： GIoU 效果也很好。
我们将实现 IoU + CIoU（使用安全数值）。
4) 基于锚点的 YOLO：网格、锚点、预测
YOLO 头部在每个网格位置进行预测。假设特征图是 S x S （例如，80×80）。每个单元格都可以预测 A 锚点（例如，3）。对于每个锚点，预测结果为：
盒子偏移:
tx, ty, tw, th对象性逻辑:
to类逻辑:
tc1..tcC
因此，每个尺度的张量形状为：(B, A*(5+C), S, S) or (B, A, S, S, 5+C) 重塑之后。
如何将偏移量变成真正的盒子
YOLO风格的常见解码方式（几种有效变体之一）：
bx = (sigmoid(tx) + cx) / Sby = (sigmoid(ty) + cy) / Sbw = (anchor_w * exp(tw)) / img_w（或按 S 归一化）bh = (anchor_h * exp(th)) / img_h
地点 (cx, cy) 是整数网格坐标。
重要提示您的编码/解码必须与目标分配的编码匹配。
5）数据集准备
注释格式
您的自定义数据集可以是：
COCO JSON
Pascal VOC XML
YOLO txt（类别 cx cy wh 归一化）
我们将支持通用的内部表示：
每个样本返回：
image: Tensor [3, H, W]targets: Tensor [N, 6]包含列：[class, x1, y1, x2, y2, image_index(optional)]
增强
对于目标检测，数据增强也必须对边界框进行变换：
调整大小/信箱
随机水平翻转
颜色抖动
随机仿射（可选）
马赛克/混合（高级；可选）
为了使本指南易于实施且不影响几何结构的稳定性，我们将这样做：
调整大小/信箱模式
随机翻转
HSV抖动（可选）
6) 构建模块：Conv-BN-Act、残差、颈部
一个干净的基线模块：
Conv2d -> BatchNorm2d -> SiLU
SiLU（又名 Swish）在 YOLOv5 类算法家族中很常见；LeakyReLU 在 YOLOv3 中很常见。
我们可以选择添加残余块以增强骨架，但即使是较小的骨架也可以用于验证管道。
7）模型设计
典型的结构：
骨干：以多个步长（8、16、32）提取特征图
颈部：结合了功能（FPN/PAN）
校长：预测每个尺度的检测输出。
我们将实现一个轻量级主干网络，生成 3 个特征图和一个简单的类似 FPN 的颈部。
8）解码预测
推断：
按尺度调整输出
(B, A, S, S, 5+C)将 sigmoid 函数应用于中心偏移量 + 对象性（以及通常的类概率）
转换为像素坐标系中的XYXY
将所有尺度合并成一个候选框列表。
按置信度阈值筛选
按类应用 NMS（或与类无关的 NMS）
9）目标分配（将GT与锚点匹配）
这就是基于锚点的YOLO（You Only Live Once，你只活一次）的核心。
对于每个真实值框：
确定应该使用哪个（哪些）比例尺来处理它（根据尺寸/锚点匹配）。
对于选定的尺度，计算 GT 框大小与每个锚点大小之间的 IoU（在该尺度的坐标系中）。
选择最佳锚点（或前 k 个锚点）。
从 GT 中心计算网格单元索引。
填充目标张量
[anchor, gy, gx]使用：盒回归目标
客观性 = 1
目标类
编码回归目标
如果使用解码：
bx = (sigmoid(tx) + cx)/S
然后瞄准txissigmoid^-1(bx*S - cx)但这很麻烦。
相反，YOLO式的训练往往 直接监督:
tx_target = bx*S - cx（一个 [0,1] 范围内的值）并使用 sigmoid 输出进行 BCE 训练，或使用原始输出进行 MSE 训练。tw_target = log(bw / anchor_w)（以像素或归一化单位表示）
我们将实现一个稳定版本：
预测
pxy = sigmoid(tx,ty)并使用 BCE/MSE 监督 pxy 以匹配分数偏移量预测
pwh = exp(tw,th)*anchor并使用 CIoU 对解码后的盒子进行监督（推荐）
这样就简单多了： 对解码后的盒子进行回归损失，不是直接在推特/周四上。
10）损失函数
YOLO式的亏损通常有以下特点：
盒子丢失：负责位置的预测框与 GT 框之间的 CIoU/GIoU
物体性丧失：BCEWithLogits 在对象性逻辑上
班级损失：BCEWithLogits（多标签）或 CE（单标签）
对于单标签分类（每个对象一个类别），两种方法都可行：
使用独热目标的 BCEWithLogits（YOLO 中常见）
对正位置的类别logits进行交叉熵损失（也很好）
为了保持一致性，我们将同时使用 BCEWithLogits 来处理对象和类。
处理负面情绪
你会遇到更多负面（无对象）立场。你可以：
对负客观性使用较低的权重
或者应用焦点损失（可选）
我们将实施：
具有正负权重的客观性损失。
11）训练循环
稳定性/性能的关键特性：
AMP (torch.cuda.amp)
渐变剪裁 （可选）
EMA 权重（可选，但有帮助）
LR调度器（余弦或步长）
前几个周期/步骤的热身
在 Playwright 的屏幕截图上使用此方法：
12) NMS
非极大值抑制会移除重叠的重复项。典型步骤：
按置信度对箱子进行排序
遍历最高置信区间，抑制 IoU > 阈值的框
使用基于类别的非极大值抑制进行多类别检测。
13) mAP 评估
平均精度均值要求：
对于每个类别，计算 IoU 阈值下的精确率-召回率曲线。
计算曲线下面积（AP）
各类别平均值（mAP）
COCO 使用 IoU 阈值 0.50 到 0.95 之间，步长为 0.05 的 mAP。
我们将实施：
平均AP@0.5
以及可选的 COCO 风格的 mAP@[.5:.95]
14）可视化
在正式训练之前，先进行如下的想象：
每个规模的目标分配
经过几次迭代后解码的预测结果
NMS 输出
这样可以解决 80% 的“我的模型无法学习”的问题。
15) 完整核心实现（参考代码）
以下是一套精简但完整的核心文件，您可以将其放入代码仓库中。它虽然不算“小巧”，但易于阅读且设计合理，确保了代码的正确性。
15.1 仓库结构
yolo_scratch/
README.md
train.py
eval.py
predict.py
yolo/
__init__.py
model.py
modules.py
loss.py
assigner.py
box_ops.py
nms.py
metrics.py
data.py
transforms.py
utils.py
configs/
coco.yaml
custom.yaml
yolo_scratch/
README.md
train.py
eval.py
predict.py
yolo/
__init__.py
model.py
modules.py
loss.py
assigner.py
box_ops.py
nms.py
metrics.py
data.py
transforms.py
utils.py
configs/
coco.yaml
custom.yaml
15.2 yolo/box_ops.py
import torch
def xyxy_to_xywh(boxes: torch.Tensor) -> torch.Tensor:
# boxes: [..., 4]
x1, y1, x2, y2 = boxes.unbind(-1)
cx = (x1 + x2) * 0.5
cy = (y1 + y2) * 0.5
w = (x2 - x1).clamp(min=0)
h = (y2 - y1).clamp(min=0)
return torch.stack([cx, cy, w, h], dim=-1)
def xywh_to_xyxy(boxes: torch.Tensor) -> torch.Tensor:
cx, cy, w, h = boxes.unbind(-1)
half_w = w * 0.5
half_h = h * 0.5
x1 = cx - half_w
y1 = cy - half_h
x2 = cx + half_w
y2 = cy + half_h
return torch.stack([x1, y1, x2, y2], dim=-1)
def box_iou_xyxy(boxes1: torch.Tensor, boxes2: torch.Tensor, eps: float = 1e-9) -> torch.Tensor:
# boxes1: [N,4], boxes2: [M,4]
x11, y11, x12, y12 = boxes1[:, 0], boxes1[:, 1], boxes1[:, 2], boxes1[:, 3]
x21, y21, x22, y22 = boxes2[:, 0], boxes2[:, 1], boxes2[:, 2], boxes2[:, 3]
inter_x1 = torch.maximum(x11[:, None], x21[None, :])
inter_y1 = torch.maximum(y11[:, None], y21[None, :])
inter_x2 = torch.minimum(x12[:, None], x22[None, :])
inter_y2 = torch.minimum(y12[:, None], y22[None, :])
inter_w = (inter_x2 - inter_x1).clamp(min=0)
inter_h = (inter_y2 - inter_y1).clamp(min=0)
inter = inter_w * inter_h
area1 = (x12 - x11).clamp(min=0) * (y12 - y11).clamp(min=0)
area2 = (x22 - x21).clamp(min=0) * (y22 - y21).clamp(min=0)
union = area1[:, None] + area2[None, :] - inter
return inter / (union + eps)
def ciou_loss_xyxy(pred: torch.Tensor, target: torch.Tensor, eps: float = 1e-7) -> torch.Tensor:
"""
pred, target: [N,4] in xyxy
Returns: [N] CIoU loss = 1 - CIoU
"""
# IoU
iou = box_iou_xyxy(pred, target).diag() # [N]
# centers and sizes
p = xyxy_to_xywh(pred)
t = xyxy_to_xywh(target)
pcx, pcy, pw, ph = p.unbind(-1)
tcx, tcy, tw, th = t.unbind(-1)
# center distance
center_dist2 = (pcx - tcx) ** 2 + (pcy - tcy) ** 2
# smallest enclosing box diagonal squared
x1 = torch.minimum(pred[:, 0], target[:, 0])
y1 = torch.minimum(pred[:, 1], target[:, 1])
x2 = torch.maximum(pred[:, 2], target[:, 2])
y2 = torch.maximum(pred[:, 3], target[:, 3])
c2 = (x2 - x1) ** 2 + (y2 - y1) ** 2 + eps
diou = iou - center_dist2 / c2
# aspect ratio penalty
v = (4 / (torch.pi ** 2)) * (torch.atan(tw / (th + eps)) - torch.atan(pw / (ph + eps))) ** 2
with torch.no_grad():
alpha = v / (1 - iou + v + eps)
ciou = diou - alpha * v
return 1 - ciou.clamp(min=-1.0, max=1.0)
15.3 yolo/nms.py
import torch
from .box_ops import box_iou_xyxy
def nms_xyxy(boxes: torch.Tensor, scores: torch.Tensor, iou_thresh: float = 0.5) -> torch.Tensor:
"""
boxes: [N,4], scores: [N]
returns indices kept
"""
if boxes.numel() == 0:
return torch.empty((0,), dtype=torch.long, device=boxes.device)
idxs = scores.argsort(descending=True)
keep = []
while idxs.numel() > 0:
i = idxs[0]
keep.append(i)
if idxs.numel() == 1:
break
rest = idxs[1:]
ious = box_iou_xyxy(boxes[i].unsqueeze(0), boxes[rest]).squeeze(0)
idxs = rest[ious <= iou_thresh]
return torch.stack(keep)
def batched_nms_xyxy(boxes, scores, labels, iou_thresh=0.5):
"""
Class-wise NMS by offsetting boxes or by filtering per class.
Here: filter per class (clear and correct).
"""
keep_all = []
for c in labels.unique():
mask = labels == c
keep = nms_xyxy(boxes[mask], scores[mask], iou_thresh)
keep_all.append(mask.nonzero(as_tuple=False).squeeze(1)[keep])
if not keep_all:
return torch.empty((0,), dtype=torch.long, device=boxes.device)
return torch.cat(keep_all)
15.4 yolo/modules.py
import torch
import torch.nn as nn
class ConvBNAct(nn.Module):
def __init__(self, in_ch, out_ch, k=3, s=1, p=None, act=True):
super().__init__()
if p is None:
p = k // 2
self.conv = nn.Conv2d(in_ch, out_ch, k, s, p, bias=False)
self.bn = nn.BatchNorm2d(out_ch)
self.act = nn.SiLU(inplace=True) if act else nn.Identity()
def forward(self, x):
return self.act(self.bn(self.conv(x)))
class Residual(nn.Module):
def __init__(self, ch):
super().__init__()
self.block = nn.Sequential(
ConvBNAct(ch, ch, 1, 1),
ConvBNAct(ch, ch, 3, 1),
)
def forward(self, x):
return x + self.block(x)
class CSPBlock(nn.Module):
"""
Light CSP-like block: split channels, apply residuals on one branch, then concat.
"""
def __init__(self, ch, n=1):
super().__init__()
c_ = ch // 2
self.conv1 = ConvBNAct(ch, c_, 1, 1)
self.conv2 = ConvBNAct(ch, c_, 1, 1)
self.m = nn.Sequential(*[Residual(c_) for _ in range(n)])
self.conv3 = ConvBNAct(2 * c_, ch, 1, 1)
def forward(self, x):
y1 = self.m(self.conv1(x))
y2 = self.conv2(x)
return self.conv3(torch.cat([y1, y2], dim=1))
15.5 yolo/model.py
import torch
import torch.nn as nn
from .modules import ConvBNAct, CSPBlock
class TinyBackbone(nn.Module):
"""
Produces 3 feature maps at strides 8, 16, 32.
"""
def __init__(self, in_ch=3, base=32):
super().__init__()
self.stem = nn.Sequential(
ConvBNAct(in_ch, base, 3, 2), # stride 2
ConvBNAct(base, base*2, 3, 2), # stride 4
CSPBlock(base*2, n=1),
)
self.stage3 = nn.Sequential(
ConvBNAct(base*2, base*4, 3, 2), # stride 8
CSPBlock(base*4, n=2),
)
self.stage4 = nn.Sequential(
ConvBNAct(base*4, base*8, 3, 2), # stride 16
CSPBlock(base*8, n=2),
)
self.stage5 = nn.Sequential(
ConvBNAct(base*8, base*16, 3, 2), # stride 32
CSPBlock(base*16, n=1),
)
def forward(self, x):
x = self.stem(x)
p3 = self.stage3(x)
p4 = self.stage4(p3)
p5 = self.stage5(p4)
return p3, p4, p5
class SimpleFPN(nn.Module):
def __init__(self, ch3, ch4, ch5, out_ch=128):
super().__init__()
self.lat5 = ConvBNAct(ch5, out_ch, 1, 1)
self.lat4 = ConvBNAct(ch4, out_ch, 1, 1)
self.lat3 = ConvBNAct(ch3, out_ch, 1, 1)
self.out4 = ConvBNAct(out_ch, out_ch, 3, 1)
self.out3 = ConvBNAct(out_ch, out_ch, 3, 1)
def forward(self, p3, p4, p5):
p5 = self.lat5(p5)
p4 = self.lat4(p4) + torch.nn.functional.interpolate(p5, scale_factor=2, mode="nearest")
p3 = self.lat3(p3) + torch.nn.functional.interpolate(p4, scale_factor=2, mode="nearest")
p4 = self.out4(p4)
p3 = self.out3(p3)
return p3, p4, p5
class DetectHead(nn.Module):
def __init__(self, in_ch, num_anchors, num_classes):
super().__init__()
self.num_anchors = num_anchors
self.num_classes = num_classes
self.pred = nn.Conv2d(in_ch, num_anchors * (5 + num_classes), 1, 1, 0)
def forward(self, x):
return self.pred(x)
class YOLO(nn.Module):
def __init__(self, num_classes, anchors, base=32):
"""
anchors: list of 3 scales, each is list of (w,h) in pixels for the model input size (e.g., 640)
e.g. [
[(10,13),(16,30),(33,23)], # stride 8
[(30,61),(62,45),(59,119)], # stride 16
[(116,90),(156,198),(373,326)] # stride 32
]
"""
super().__init__()
self.num_classes = num_classes
self.anchors = anchors
self.backbone = TinyBackbone(in_ch=3, base=base)
# backbone channels: p3=base*4, p4=base*8, p5=base*16
self.fpn = SimpleFPN(base*4, base*8, base*16, out_ch=base*4)
na = len(anchors[0])
self.head3 = DetectHead(base*4, na, num_classes)
self.head4 = DetectHead(base*4, na, num_classes)
self.head5 = DetectHead(base*4, na, num_classes)
def forward(self, x):
p3, p4, p5 = self.backbone(x)
f3, f4, f5 = self.fpn(p3, p4, p5)
o3 = self.head3(f3)
o4 = self.head4(f4)
o5 = self.head5(f5)
return [o3, o4, o5]
15.6 yolo/assigner.py （目标分配）
import torch
def build_targets(
targets, # list of length B, each: Tensor [Ni, 5] -> (cls, x1, y1, x2, y2) in pixels
anchors, # per scale: list of (w,h) in pixels at input size
strides, # [8,16,32]
img_size, # int, e.g. 640
num_classes,
device
):
"""
Returns per-scale target tensors:
tbox: list of [B, A, S, S, 4] in xyxy pixels
tobj: list of [B, A, S, S] (0/1)
tcls: list of [B, A, S, S, C] one-hot
indices: list of tuples for positives (b, a, gy, gx)
"""
B = len(targets)
out = []
indices_all = []
for scale_idx, (anc, stride) in enumerate(zip(anchors, strides)):
S = img_size // stride
A = len(anc)
tbox = torch.zeros((B, A, S, S, 4), device=device)
tobj = torch.zeros((B, A, S, S), device=device)
tcls = torch.zeros((B, A, S, S, num_classes), device=device)
indices = []
anc_wh = torch.tensor(anc, device=device, dtype=torch.float32) # [A,2]
for b in range(B):
if targets[b].numel() == 0:
continue
gt = targets[b].to(device)
cls = gt[:, 0].long()
x1y1 = gt[:, 1:3]
x2y2 = gt[:, 3:5]
gxy = (x1y1 + x2y2) * 0.5
gwh = (x2y2 - x1y1).clamp(min=1.0)
# pick best anchor by IoU of width/height (approx)
# IoU(wh) = min(w)/max(w) * min(h)/max(h)
wh = gwh[:, None, :] # [N,1,2]
min_wh = torch.minimum(wh, anc_wh[None, :, :])
max_wh = torch.maximum(wh, anc_wh[None, :, :])
iou_wh = (min_wh[..., 0] / max_wh[..., 0]) * (min_wh[..., 1] / max_wh[..., 1]) # [N,A]
best_a = torch.argmax(iou_wh, dim=1) # [N]
# grid cell
gx = (gxy[:, 0] / stride).clamp(min=0, max=S-1e-3)
gy = (gxy[:, 1] / stride).clamp(min=0, max=S-1e-3)
gi = gx.long()
gj = gy.long()
for i in range(gt.shape[0]):
a = best_a[i].item()
x1, y1, x2, y2 = gt[i, 1:].tolist()
# assign
tobj[b, a, gj[i], gi[i]] = 1.0
tbox[b, a, gj[i], gi[i]] = torch.tensor([x1, y1, x2, y2], device=device)
tcls[b, a, gj[i], gi[i], cls[i]] = 1.0
indices.append((b, a, gj[i].item(), gi[i].item()))
out.append((tbox, tobj, tcls))
indices_all.append(indices)
return out, indices_all
15.7 yolo/loss.py
import torch
import torch.nn as nn
from .box_ops import xywh_to_xyxy, ciou_loss_xyxy
class YOLOLoss(nn.Module):
def __init__(self, anchors, strides, num_classes, img_size,
lambda_box=7.5, lambda_obj=1.0, lambda_cls=1.0,
obj_pos_weight=1.0, obj_neg_weight=0.5):
super().__init__()
self.anchors = anchors
self.strides = strides
self.num_classes = num_classes
self.img_size = img_size
self.lambda_box = lambda_box
self.lambda_obj = lambda_obj
self.lambda_cls = lambda_cls
self.bce = nn.BCEWithLogitsLoss(reduction="none")
self.obj_pos_weight = obj_pos_weight
self.obj_neg_weight = obj_neg_weight
def decode_scale(self, pred, scale_idx):
"""
pred: [B, A*(5+C), S, S]
returns:
boxes_xyxy: [B, A, S, S, 4] in pixels
obj_logit: [B, A, S, S]
cls_logit: [B, A, S, S, C]
"""
B, _, S, _ = pred.shape
A = len(self.anchors[scale_idx])
C = self.num_classes
stride = self.strides[scale_idx]
pred = pred.view(B, A, 5 + C, S, S).permute(0, 1, 3, 4, 2).contiguous()
# [B, A, S, S, 5+C]
tx_ty = pred[..., 0:2]
tw_th = pred[..., 2:4]
obj = pred[..., 4]
cls = pred[..., 5:]
# grid
gy, gx = torch.meshgrid(torch.arange(S, device=pred.device),
torch.arange(S, device=pred.device), indexing="ij")
grid = torch.stack([gx, gy], dim=-1).float() # [S,S,2]
# anchors
anc = torch.tensor(self.anchors[scale_idx], device=pred.device).float() # [A,2]
anc = anc.view(1, A, 1, 1, 2)
# decode center
pxy = (tx_ty.sigmoid() + grid.view(1, 1, S, S, 2)) * stride # pixels
# decode wh
pwh = (tw_th.exp() * anc) # pixels
boxes_xywh = torch.cat([pxy, pwh], dim=-1)
boxes_xyxy = xywh_to_xyxy(boxes_xywh)
return boxes_xyxy, obj, cls
def forward(self, preds, targets_per_scale):
"""
preds: list of 3 scale outputs
targets_per_scale: list of (tbox, tobj, tcls) per scale
"""
total_box = torch.tensor(0.0, device=preds[0].device)
total_obj = torch.tensor(0.0, device=preds[0].device)
total_cls = torch.tensor(0.0, device=preds[0].device)
for s, pred in enumerate(preds):
tbox, tobj, tcls = targets_per_scale[s]
pbox, pobj_logit, pcls_logit = self.decode_scale(pred, s)
# objectness loss with weighting
obj_loss = self.bce(pobj_logit, tobj)
w = torch.where(tobj > 0.5,
torch.full_like(obj_loss, self.obj_pos_weight),
torch.full_like(obj_loss, self.obj_neg_weight))
total_obj = total_obj + (obj_loss * w).mean()
# positives mask
pos = tobj > 0.5
if pos.any():
# box loss CIoU
pbox_pos = pbox[pos]
tbox_pos = tbox[pos]
box_loss = ciou_loss_xyxy(pbox_pos, tbox_pos).mean()
total_box = total_box + box_loss
# class loss
cls_loss = self.bce(pcls_logit[pos], tcls[pos]).mean()
total_cls = total_cls + cls_loss
loss = self.lambda_box * total_box + self.lambda_obj * total_obj + self.lambda_cls * total_cls
return loss, {"box": total_box.detach(), "obj": total_obj.detach(), "cls": total_cls.detach()}
15.8 yolo/utils.py
import torch
def set_seed(seed=42):
import random, numpy as np
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
class AverageMeter:
def __init__(self):
self.sum = 0.0
self.count = 0
def update(self, v, n=1):
self.sum += float(v) * n
self.count += n
@property
def avg(self):
return self.sum / max(1, self.count)
15.9 yolo/transforms.py （简易信箱 + 翻转）
import torch
import torchvision.transforms.functional as TF
def letterbox(image, boxes_xyxy, new_size=640):
"""
image: PIL or Tensor [C,H,W]
boxes_xyxy: Tensor [N,4] in pixels
returns resized/padded image and transformed boxes
"""
if not torch.is_tensor(image):
image = TF.to_tensor(image)
c, h, w = image.shape
scale = min(new_size / h, new_size / w)
nh, nw = int(round(h * scale)), int(round(w * scale))
image_resized = TF.resize(image, [nh, nw])
pad_h = new_size - nh
pad_w = new_size - nw
top = pad_h // 2
left = pad_w // 2
image_padded = torch.zeros((c, new_size, new_size), dtype=image.dtype)
image_padded[:, top:top+nh, left:left+nw] = image_resized
if boxes_xyxy.numel() > 0:
boxes = boxes_xyxy.clone()
boxes *= scale
boxes[:, [0, 2]] += left
boxes[:, [1, 3]] += top
else:
boxes = boxes_xyxy
return image_padded, boxes
def random_hflip(image, boxes_xyxy, p=0.5):
if torch.rand(()) > p:
return image, boxes_xyxy
c, h, w = image.shape
image = torch.flip(image, dims=[2]) # flip width
boxes = boxes_xyxy.clone()
if boxes.numel() > 0:
x1 = boxes[:, 0].clone()
x2 = boxes[:, 2].clone()
boxes[:, 0] = (w - x2)
boxes[:, 2] = (w - x1)
return image, boxes
15.10 yolo/data.py （自定义数据集框架）
import os
import torch
from torch.utils.data import Dataset
from PIL import Image
from .transforms import letterbox, random_hflip
class DetectionDataset(Dataset):
"""
Expects a list of samples where each sample has:
- image_path
- annotations: list of [cls, x1, y1, x2, y2] in pixels
You can write adapters to load COCO or YOLO txt into this format.
"""
def __init__(self, samples, img_size=640, augment=True):
self.samples = samples
self.img_size = img_size
self.augment = augment
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
s = self.samples[idx]
img = Image.open(s["image_path"]).convert("RGB")
ann = s.get("annotations", [])
if len(ann) > 0:
target = torch.tensor(ann, dtype=torch.float32) # [N,5]
else:
target = torch.zeros((0,5), dtype=torch.float32)
cls = target[:, 0:1]
boxes = target[:, 1:5]
img, boxes = letterbox(img, boxes, self.img_size)
if self.augment:
img, boxes = random_hflip(img, boxes, p=0.5)
# normalize image
img = img.clamp(0, 1)
# pack back
if boxes.numel() > 0:
target = torch.cat([cls, boxes], dim=1)
else:
target = torch.zeros((0,5), dtype=torch.float32)
return img, target
def collate_fn(batch):
images, targets = zip(*batch)
images = torch.stack(images, dim=0)
# targets remains list[Tensor]
return images, list(targets)
15.11 train.py （端到端训练循环）
import torch
from torch.utils.data import DataLoader
from yolo.model import YOLO
from yolo.loss import YOLOLoss
from yolo.assigner import build_targets
from yolo.utils import set_seed, AverageMeter
from yolo.data import DetectionDataset, collate_fn
def train_one_epoch(model, loss_fn, loader, optimizer, device, anchors, strides, img_size, num_classes, scaler=None):
model.train()
meter = AverageMeter()
for images, targets_list in loader:
images = images.to(device)
# build targets per scale
targets_per_scale, _ = build_targets(
targets_list, anchors=anchors, strides=strides,
img_size=img_size, num_classes=num_classes, device=device
)
targets_per_scale = [(tbox, tobj, tcls) for (tbox, tobj, tcls) in targets_per_scale]
optimizer.zero_grad(set_to_none=True)
if scaler is not None:
with torch.cuda.amp.autocast():
preds = model(images)
loss, logs = loss_fn(preds, targets_per_scale)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
preds = model(images)
loss, logs = loss_fn(preds, targets_per_scale)
loss.backward()
optimizer.step()
meter.update(loss.item(), n=images.size(0))
return meter.avg
def main():
set_seed(42)
device = "cuda" if torch.cuda.is_available() else "cpu"
img_size = 640
num_classes = 80 # COCO example
strides = [8, 16, 32]
anchors = [
[(10,13),(16,30),(33,23)],
[(30,61),(62,45),(59,119)],
[(116,90),(156,198),(373,326)]
]
# TODO: load your samples list here
samples = [] # [{"image_path": "...", "annotations": [[cls,x1,y1,x2,y2], ...]}, ...]
ds = DetectionDataset(samples, img_size=img_size, augment=True)
loader = DataLoader(ds, batch_size=8, shuffle=True, num_workers=4, pin_memory=True, collate_fn=collate_fn)
model = YOLO(num_classes=num_classes, anchors=anchors, base=32).to(device)
loss_fn = YOLOLoss(anchors, strides, num_classes, img_size).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4, weight_decay=1e-4)
scaler = torch.cuda.amp.GradScaler() if device == "cuda" else None
for epoch in range(1, 101):
avg_loss = train_one_epoch(model, loss_fn, loader, optimizer, device, anchors, strides, img_size, num_classes, scaler)
print(f"Epoch {epoch:03d} | loss={avg_loss:.4f}")
torch.save(model.state_dict(), "yolo_scratch.pt")
if __name__ == "__main__":
main()
16）推理：解码 + NMS（预测流程）
16.1 解码助手（添加到 yolo/metrics.py or yolo/utils.py)
import torch
from .box_ops import xywh_to_xyxy
from .nms import batched_nms_xyxy
@torch.no_grad()
def decode_predictions(preds, anchors, strides, num_classes, conf_thresh=0.25, iou_thresh=0.5):
"""
preds: list of 3 tensors [B, A*(5+C), S, S]
returns per image: boxes [M,4], scores [M], labels [M]
"""
outputs = []
for b in range(preds[0].shape[0]):
boxes_all = []
scores_all = []
labels_all = []
for s, p in enumerate(preds):
B, _, S, _ = p.shape
A = len(anchors[s])
C = num_classes
stride = strides[s]
x = p[b:b+1].view(1, A, 5+C, S, S).permute(0,1,3,4,2).contiguous()[0] # [A,S,S,5+C]
tx_ty = x[..., 0:2]
tw_th = x[..., 2:4]
obj_logit = x[..., 4]
cls_logit = x[..., 5:]
gy, gx = torch.meshgrid(torch.arange(S, device=p.device),
torch.arange(S, device=p.device), indexing="ij")
grid = torch.stack([gx, gy], dim=-1).float() # [S,S,2]
anc = torch.tensor(anchors[s], device=p.device).float().view(A,1,1,2)
pxy = (tx_ty.sigmoid() + grid) * stride
pwh = tw_th.exp() * anc
boxes_xywh = torch.cat([pxy, pwh], dim=-1) # [A,S,S,4]
boxes_xyxy = xywh_to_xyxy(boxes_xywh)
obj = obj_logit.sigmoid() # [A,S,S]
cls_prob = cls_logit.sigmoid() # [A,S,S,C]
# combine: per-class confidence = obj * cls_prob
conf = obj.unsqueeze(-1) * cls_prob # [A,S,S,C]
conf = conf.view(-1, C)
boxes = boxes_xyxy.view(-1, 4)
scores, labels = conf.max(dim=1)
keep = scores > conf_thresh
boxes_all.append(boxes[keep])
scores_all.append(scores[keep])
labels_all.append(labels[keep])
if boxes_all:
boxes = torch.cat(boxes_all, dim=0)
scores = torch.cat(scores_all, dim=0)
labels = torch.cat(labels_all, dim=0)
keep = batched_nms_xyxy(boxes, scores, labels, iou_thresh=iou_thresh)
outputs.append((boxes[keep], scores[keep], labels[keep]))
else:
outputs.append((torch.zeros((0,4)), torch.zeros((0,)), torch.zeros((0,), dtype=torch.long)))
return outputs
17) mAP 评估器（核心逻辑）
一个正确的 mAP 实现程序比较冗长；这里提供一种简洁、最小化的评估器方法：
收集每张图像的预测结果：方框、分数、标签
每张图片收集GT值：盒子、标签
针对每个班级：
按得分对预测结果进行排序
使用高于阈值的最佳 IoU 匹配来标记 TP/FP（并且每个 GT 只匹配一次）
计算精确率-召回率曲线
利用数值积分法计算等差数列。
各类别平均值 → mAP
如果您想要 COCO 风格的 mAP@[.5:.95]，请在多个 IoU 阈值下重复上述操作并取平均值。
如果你需要，我可以粘贴一个完整的、可以直接运行的程序。
metrics.py我将 mAP@0.5 和 COCO mAP 合并到一个文件中（文件有几百行）。为了便于阅读，本文主要侧重于 YOLO 流程。
18) COCO 培训（实践笔记）
为了有效地进行 COCO 训练：
如果可以，请使用更坚固的骨架和更大的批次。
使用多尺度训练（每批次随机改变输入大小）
在前1-3个epoch中使用预热阶段进行LR训练
使用EMA权重进行评估
注意：
目标丢失爆炸（通常是目标分配错误）
接近于零的正向信号（锚点/步幅不匹配）
方框漂移到图像外（解码错误）
19) 使用自定义数据集进行训练（最快正确的工作流程）
选择模型输入大小：基线模型通常为 640。
将注释转换为绝对 XYXY 像素
在训练前可视化图像上的方框
从...开始：
没有花哨的增强
小型号
对 20 张图片进行过拟合
如果模型过拟合，则扩大规模：
更多数据
扩增
更好的脊柱/颈部
20）常见错误检查清单（节省时间）
方框格式错误 （XYWH vs XYXY）
方框已归一化，但仍被视为像素 （或相反亦然）
目标被分配到错误的尺度 （步幅不匹配）
锚点大小与输入大小不匹配 （锚点为 416，但训练点为 640）
交换 x/y 索引 （张量索引中的 gx 与 gy）
忘记限制网格索引
在转换为 XYXY 之前应用 NMS
mAP评估将多个预测结果与同一个GT进行匹配。
在 PyTorch 中从头开始实现 YOLO 是真正理解现代目标检测的最佳方法之一——因为你被迫将每个移动部分连接起来：标签如何变成训练目标，预测如何变成真正的框，为什么锚框存在，以及目标性实际上学习的是什么。
完成本次构建后，您应该拥有一个完整且可正常运行的探测器，其功能包括：
一个多尺度的YOLO风格模型（脊柱+颈部+检测头）
正确的目标分配流程（真实值 → 网格单元 + 锚框）
稳定的损失函数设置（CIoU/GIoU 用于框，BCE 用于对象性和类别）
正确的推理路径（解码→置信度过滤→非极大值抑制）
一条清晰的生产级评估途径（mAP@0.5 和 COCO mAP@[.5:.95]）
一个可以扩展到实际项目中的代码仓库结构
最重要的启示是，YOLO并非“魔法”——它是一套精心设计的系统。 一致的坐标变换, 负责任的锚匹配和 平衡损失如果这些要素中任何一个不一致（例如错误的框格式、错误的步长、不匹配的锚点、x/y 坐标反转），学习过程就会崩溃。但当所有要素都匹配时，模型就能顺利训练并良好扩展。
