以 YOLOv3(yolov3.yaml)作为模型架构,COCO128(coco128.yaml)作为数据集,并采用 416×416 输入图像尺寸为例。

训练入口是train.py的train函数,另外训练还涉及到如下文件:
(1)utils/ 目录下多个文件:
dataloaders.py → 数据加载
loss.py → 损失计算
general.py → 工具函数
torch_utils.py → 设备、EMA 等
loggers/ → 日志记录
callbacks.py → 回调机制

(2)数据与配置:
data/coco128.yaml
data/hyps/hyp.scratch-low.yaml
权重文件(.pt)

其一般训练过程如下:

训练循环开始
    ↓
数据加载 → 前向传播 → 损失计算
    ↓               ↓
梯度累积 ← 反向传播 ← 损失计算
    ↓
参数更新 → 学习率调整
    ↓
日志记录 → 模型保存
    ↓
下一个epoch

具体下面展开。

1 训练准备

超参数加载

if isinstance(hyp, str):
    with open(hyp, errors="ignore") as f:
        hyp = yaml.safe_load(f)  # 加载超参数字典
LOGGER.info(colorstr("hyperparameters: ") + ", ".join(f"{k}={v}" for k, v in hyp.items()))
opt.hyp = hyp.copy()  # 保存到检查点

数据加载器创建:

# 训练数据加载器
train_loader, dataset = create_dataloader(
    train_path,
    imgsz,
    batch_size // WORLD_SIZE,
    gs,
    single_cls,
    hyp=hyp,
    augment=True,
    cache=None if opt.cache == "val" else opt.cache,
    rect=opt.rect,
    rank=LOCAL_RANK,
    workers=workers,
    image_weights=opt.image_weights,
    quad=opt.quad,
    prefix=colorstr("train: "),
    shuffle=True,
    seed=opt.seed,
)

# 验证数据加载器(仅主进程)
if RANK in {-1, 0}:
    val_loader = create_dataloader(
        val_path,
        imgsz,
        batch_size // WORLD_SIZE * 2,
        gs,
        single_cls,
        hyp=hyp,
        cache=None if noval else opt.cache,
        rect=True,
        rank=-1,
        workers=workers * 2,
        pad=0.5,
        prefix=colorstr("val: "),
    )[0]

模型加载与配置

if pretrained:
    with torch_distributed_zero_first(LOCAL_RANK):
        weights = attempt_download(weights)  # 下载权重
    ckpt = torch_load(weights, map_location="cpu")  # 加载检查点
    model = Model(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device)  # 创建模型
    csd = ckpt["model"].float().state_dict()  # FP32状态字典
    csd = intersect_dicts(csd, model.state_dict(), exclude=exclude)  # 交集
    model.load_state_dict(csd, strict=False)  # 加载
    LOGGER.info(f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}")
else:
    model = Model(cfg, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device)  # 从头开始创建
freeze = [f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0]))]  # 要冻结的层
for k, v in model.named_parameters():
    v.requires_grad = True  # 默认训练所有层
    if any(x in k for x in freeze):
        LOGGER.info(f"freezing {k}")
        v.requires_grad = False

训练组件初始化:优化器、学习率调度器、指数移动平均、损失计算器

优化器:

nbs = 64  # 名义批次大小
accumulate = max(round(nbs / batch_size), 1)  # 累积梯度步数
hyp["weight_decay"] *= batch_size * accumulate / nbs  # 调整权重衰减
optimizer = smart_optimizer(model, opt.optimizer, hyp["lr0"], hyp["momentum"], hyp["weight_decay"])

学习率调度器:

if opt.cos_lr:
    lf = one_cycle(1, hyp["lrf"], epochs)  # 余弦调度
else:
    def lf(x):
        """线性学习率调度函数"""
        return (1 - x / epochs) * (1.0 - hyp["lrf"]) + hyp["lrf"]  # 线性

scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)

指数移动平均:

ema = ModelEMA(model) if RANK in {-1, 0} else None

损失计算器:

compute_loss = ComputeLoss(model)  # 初始化损失类

2 训练过程

假设 batch_size = 4,输入 4 张 416×416 图像,单batch为例。

步骤 1:数据加载

pbar = enumerate(train_loader)
# 从数据加载器获取当前batch的数据
# imgs: [4, 3, 416, 416] - 批图像
# targets: [N, 6] - 标签,格式为(image_index, class_id, x, y, w, h)
# paths: 图像路径列表

步骤 2:前向传播(Forward Pass)

# 前向传播(使用混合精度训练)
with torch.cuda.amp.autocast(amp):
    # 模型前向传播
    pred = model(imgs)  # forward,获取三个尺度的预测
    
    # 在训练模式下,Detect层直接返回未处理的卷积输出
    # 输出是包含3个张量的列表,对应P3/P4/P5三个检测层:
    #   pred[0]: [4, 255, 52, 52]  # P3/8: 52×52网格,检测小目标
    #   pred[1]: [4, 255, 26, 26]  # P4/16: 26×26网格,检测中目标  
    #   pred[2]: [4, 255, 13, 13]  # P5/32: 13×13网格,检测大目标
    # 255 = 3 anchors × (4坐标偏移 + 1置信度 + 80类别概率)

网络处理流程:
Backbone(Darknet53):提取多尺度特征
P5: (4, 1024, 13, 13) → stride=32
P4: (4, 512, 26, 26) → stride=16
P3: (4, 256, 52, 52) → stride=8

Head(FPN-style Detection):
P5分支:卷积 → 输出 (4, 255, 13, 13)
P4分支:P5上采样 → 与P4拼接 → 卷积 → 输出 (4, 255, 26, 26)
P3分支:P4上采样 → 与P3拼接 → 卷积 → 输出 (4, 255, 52, 52)

训练与推理输出差异:
训练时:返回原始卷积输出,用于损失计算。
推理时:进行坐标解码,输出(x1, y1, x2, y2, confidence, class_id)格式。

步骤 3:损失计算(Loss Computation)

# 损失计算
loss, loss_items = compute_loss(pred, targets.to(device))

# ComputeLoss类内部的核心处理流程:
# 1. 正负样本分配(基于IoU的匹配原则):
#    - 对每个真实框,确定其中心点落在哪个网格(i,j)
#    - 在该网格的3个锚框中,选择与真实框IoU最大的作为正样本
#    - 同一尺度的其他位置作为负样本(如果IoU<ignore_thresh)
#    - 每个尺度独立进行样本分配
# 
# 2. 损失计算(在编码空间进行,无需坐标解码):
#    - 坐标损失:比较预测偏移量(tx,ty,tw,th)与真实偏移量
#    - 置信度损失:二元交叉熵,区分前景/背景
#    - 分类损失:多标签二元交叉熵,支持多标签分类
# 
# 3. 损失加权与缩放(基于文档1第176-179行):
#    总损失 = hyp["box"] × L_box + hyp["obj"] × L_obj + hyp["cls"] × L_cls
#    其中损失权重会根据检测层数、类别数、图像尺寸自动缩放:
#    - hyp["box"] *= 3 / nl  # 按检测层数缩放
#    - hyp["cls"] *= nc / 80 * 3 / nl  # 按类别数和层数缩放
#    - hyp["obj"] *= (imgsz / 640) ** 2 * 3 / nl  # 按图像大小和层数缩放

关键说明:
损失计算在网络的输出空间进行,真实标签已被编码为偏移量格式。
采用"一个真实框匹配一个最佳锚框"的稀疏分配策略。
默认损失权重来自hyp.yaml(如box=0.05, obj=1.0, cls=0.5),但会动态调整。

步骤 4:反向传播与优化

# 反向传播(带梯度缩放)
scaler.scale(loss).backward()

# 梯度累积:累计accumulate个batch的梯度后更新一次
if (i + 1) % accumulate == 0 or i == len(train_loader) - 1:
    # 1. 取消梯度缩放,准备梯度裁剪
    scaler.unscale_(optimizer)
    
    # 2. 梯度裁剪,防止梯度爆炸(最大梯度范数10.0)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)
    
    # 3. 参数更新
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()
    
    # 4. 更新指数移动平均模型(EMA)
    if ema:
        ema.update(model)

# 学习率调度(每个batch后更新)
scheduler.step()

优化器:SGD with momentum=0.937
权重衰减:0.0005
初始学习率:hyp[“lr0”](默认0.01)
学习率调度:线性或余弦退火(由opt.cos_lr控制)

步骤 5:验证(Validation)
每轮结束后,在 val 集上运行:

# 每轮训练结束后,在主进程执行验证
if RANK in {-1, 0} and not noval:
    # 使用EMA模型进行更稳定的推理
    results, maps, _ = validate.run(
        data_dict,
        batch_size=batch_size // WORLD_SIZE * 2,
        imgsz=imgsz,
        model=ema.ema,
        single_cls=single_cls,
        dataloader=val_loader,
        save_dir=save_dir,
        plots=False,
        callbacks=callbacks,
        compute_loss=None,  # 验证阶段通常不计算损失
    )
    
    # 输出COCO标准评估指标:
    # - mAP@0.5: IoU阈值为0.5时的平均精度
    # - mAP@0.5:0.95: IoU阈值从0.5到0.95的平均精度
    # - precision: 精确率
    # - recall: 召回率

使验证阶段的关键差异:
使用EMA模型而非原始模型,提高稳定性。
进行NMS后处理(IoU阈值默认0.65)。
不进行损失计算,只评估检测性能。

步骤 6:保存检查点

# 保存最新检查点
ckpt = {
    'epoch': epoch,
    'best_fitness': best_fitness,
    'model': deepcopy(de_parallel(model)).half(),  # 保存半精度模型
    'ema': deepcopy(ema.ema).half(),  # 保存EMA模型
    'updates': ema.updates,
    'optimizer': optimizer.state_dict(),
    'opt': vars(opt),
    'git': GIT_INFO,  # git信息
    'date': datetime.now().isoformat(),
}

# 保存为last.pt
torch.save(ckpt, last)

# 如果当前模型性能最佳,保存为best.pt
if best_fitness == fi:
    torch.save(ckpt, best)
    
# 定期保存检查点
if (not nosave) and (opt.save_period > 0) and (epoch % opt.save_period == 0):
    torch.save(ckpt, w / f'epoch{epoch}.pt')

关键张量维度总结(Batch=4, Img=416)
在这里插入图片描述

3 小结

以 yolov3.yaml + coco128.yaml + 416x416 为例的训练流程:
加载 128 张 COCO 图像及标签
Mosaic + HSV + Flip 增强 → 416×416 batch
Darknet53 提取多尺度特征(13/26/52)
FPN 融合后输出 3 个尺度的 255 通道预测
基于 anchor-IoU 匹配,计算复合损失
SGD 优化 + EMA 平滑
每轮验证并保存 best/last 模型
这是一个端到端、单阶段、多尺度回归的经典目标检测训练范式。

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐