YOLOV11实战：从零搭建目标检测模型（附完整PyTorch代码）
·
YOLOV11实战：从零搭建目标检测模型（附完整PyTorch代码）
如果你已经对YOLO系列耳熟能详，甚至亲手调过YOLOv5或YOLOv8，那么面对最新的YOLOV11，你可能会好奇：它到底带来了哪些实质性的改变？是更快的速度，更高的精度，还是更优雅的代码架构？更重要的是，如何能避开那些官方文档里语焉不详的“坑”，快速搭建一个属于自己的、能跑起来的检测模型？这篇文章就是为你准备的。我们不打算复述那些随处可见的“网络结构三件套”介绍，而是聚焦于实战——从环境配置、数据准备、模型构建、训练调优到最终推理，我会手把手带你走一遍完整的流程，并分享我在复现和调试过程中积累的、那些在论文和官方Repo里找不到的细节与技巧。无论你是想快速验证一个新想法，还是希望将YOLOV11部署到自己的项目中，这里都有你需要的代码和思路。
1. 环境搭建与依赖管理：打造可复现的基石
在开始任何深度学习项目之前，搭建一个稳定、可复现的环境是避免后续无数麻烦的第一步。对于YOLOV11，我们不仅需要基础的PyTorch，还要考虑CUDA版本兼容性、特定依赖库的版本锁定，以及一个便于管理的项目结构。
首先，我强烈建议使用Conda或Mamba来创建独立的Python环境。这能确保你的项目依赖不会与系统或其他项目的包发生冲突。下面是一个我常用的环境创建命令，它指定了Python 3.9（一个在稳定性和兼容性上表现不错的版本）和PyTorch 2.0+。
conda create -n yolov11 python=3.9 -y
conda activate yolov11
接下来是核心依赖的安装。PyTorch的安装需要根据你的CUDA版本（通过 nvidia-smi 查看）到其官网获取对应的安装命令。这里以CUDA 11.8为例：
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
除了PyTorch，我们还需要一系列辅助库来处理数据、可视化结果和评估模型。我习惯创建一个 requirements.txt 文件来管理它们，这比一条长长的 pip install 命令更清晰，也便于团队协作和后续部署。
# requirements.txt
opencv-python>=4.8.0 # 用于图像读取和基础处理
numpy>=1.24.0 # 数值计算基础
matplotlib>=3.7.0 # 结果可视化
pillow>=10.0.0 # 图像处理
scipy>=1.11.0 # 用于计算IoU等指标
pycocotools>=2.0.7 # COCO数据集评估标准（即使你不用COCO，很多评估函数也依赖它）
tqdm>=4.66.0 # 进度条，让训练过程更友好
tensorboard>=2.14.0 # 训练过程可视化，强烈推荐
albumentations>=1.3.0 # 功能强大的数据增强库
使用 pip install -r requirements.txt 一键安装。这里特别提一下 Albumentations，它比PyTorch自带的 transforms 提供了更多样、更高效的数据增强操作，对于提升模型鲁棒性至关重要。
注意：依赖版本冲突是深度学习项目中最令人头疼的问题之一。如果你在后续步骤中遇到奇怪的报错，首先检查各个库的版本是否兼容。一个可行的策略是，在项目根目录下使用
pip freeze > requirements_lock.txt导出当前成功运行环境的所有精确版本，作为最终的“快照”。
最后，规划一下你的项目目录。一个清晰的结构能让代码维护和实验管理事半功倍。我推荐的结构如下：
yolov11_project/
├── data/
│ ├── coco/ # 或你的自定义数据集
│ │ ├── images/
│ │ └── labels/
│ └── dataset.yaml # 数据集配置文件
├── models/
│ ├── __init__.py
│ ├── backbone.py # 骨干网络定义
│ ├── neck.py # 颈部网络定义
│ ├── head.py # 检测头定义
│ └── yolo.py # 整体模型组装
├── utils/
│ ├── datasets.py # 数据加载与预处理
│ ├── losses.py # 损失函数
│ ├── metrics.py # 评估指标
│ └── general.py # 通用函数（如NMS）
├── config/
│ └── train.yaml # 训练超参数配置
├── runs/ # TensorBoard日志、权重保存
├── train.py # 训练脚本
├── detect.py # 推理脚本
└── evaluate.py # 评估脚本
2. 数据准备与高效加载：告别混乱的起点
模型的表现，七分靠数据。对于目标检测任务，数据的格式和组织方式直接决定了后续代码的复杂度。YOLO系列通常使用一种简单的文本标注格式：每个图像对应一个 .txt 文件，每行代表一个目标，格式为 class_id x_center y_center width height，其中坐标和宽高都是相对于图像宽度和高度的归一化值（范围0-1）。
假设我们处理的是经典的COCO数据集，或者一个自建的数据集，我们需要将其转换为这种格式。下面是一个简单的Python脚本片段，用于将COCO的JSON格式标注转换为YOLO格式：
import json
from pathlib import Path
def convert_coco_to_yolo(coco_json_path, output_label_dir):
with open(coco_json_path, 'r') as f:
coco_data = json.load(f)
# 创建类别ID到连续ID的映射（YOLO通常从0开始）
categories = {cat['id']: idx for idx, cat in enumerate(coco_data['categories'])}
# 按图像ID组织标注
from collections import defaultdict
img_anns = defaultdict(list)
for ann in coco_data['annotations']:
img_anns[ann['image_id']].append(ann)
for img in coco_data['images']:
img_id = img['id']
img_w, img_h = img['width'], img['height']
label_lines = []
for ann in img_anns.get(img_id, []):
# COCO标注是[x_top_left, y_top_left, width, height]
x_tl, y_tl, w, h = ann['bbox']
# 转换为YOLO格式：中心点(x_center, y_center)和归一化的宽高
x_center = (x_tl + w / 2) / img_w
y_center = (y_tl + h / 2) / img_h
w_norm = w / img_w
h_norm = h / img_h
class_id = categories[ann['category_id']]
label_lines.append(f"{class_id} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}")
# 写入标签文件
label_path = Path(output_label_dir) / f"{Path(img['file_name']).stem}.txt"
label_path.parent.mkdir(parents=True, exist_ok=True)
if label_lines:
label_path.write_text('\n'.join(label_lines))
else: # 没有目标的图像，写一个空文件
label_path.write_text('')
数据准备好后，我们需要一个高效的 Dataset 类来加载它们。这里的关键在于数据增强。我倾向于使用Albumentations库来构建一个强大的增强流水线，它支持同时处理图像和边界框，并且速度很快。
import cv2
import torch
from torch.utils.data import Dataset
import albumentations as A
from albumentations.pytorch import ToTensorV2
class YOLODataset(Dataset):
def __init__(self, img_dir, label_dir, img_size=640, augment=True):
self.img_dir = Path(img_dir)
self.label_dir = Path(label_dir)
self.img_size = img_size
self.augment = augment
# 收集所有图像路径
self.img_paths = list(self.img_dir.glob('*.jpg')) + list(self.img_dir.glob('*.png'))
self.img_paths.sort()
# 定义训练和验证时的数据增强
if self.augment:
self.transform = A.Compose([
A.RandomResizedCrop(height=img_size, width=img_size, scale=(0.5, 1.0)),
A.HorizontalFlip(p=0.5),
A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.5),
A.Blur(blur_limit=3, p=0.1),
A.ToGray(p=0.05),
A.Normalize(mean=[0, 0, 0], std=[1, 1, 1]), # 先归一化到0-1，训练时再标准化
ToTensorV2(),
], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels'], min_visibility=0.3))
else:
self.transform = A.Compose([
A.LongestMaxSize(max_size=img_size),
A.PadIfNeeded(min_height=img_size, min_width=img_size, border_mode=cv2.BORDER_CONSTANT, value=0),
A.Normalize(mean=[0, 0, 0], std=[1, 1, 1]),
ToTensorV2(),
], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))
def __len__(self):
return len(self.img_paths)
def __getitem__(self, idx):
img_path = self.img_paths[idx]
image = cv2.imread(str(img_path))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # OpenCV默认是BGR
label_path = self.label_dir / f"{img_path.stem}.txt"
boxes = []
class_labels = []
if label_path.exists():
with open(label_path, 'r') as f:
for line in f.readlines():
parts = line.strip().split()
if len(parts) == 5:
class_id, xc, yc, w, h = map(float, parts)
boxes.append([xc, yc, w, h])
class_labels.append(class_id)
if self.augment and len(boxes) > 0:
transformed = self.transform(image=image, bboxes=boxes, class_labels=class_labels)
image = transformed['image']
boxes = transformed['bboxes']
class_labels = transformed['class_labels']
else:
transformed = self.transform(image=image, bboxes=boxes, class_labels=class_labels)
image = transformed['image']
# 验证时，增强不改变框的位置，但需要确保格式一致
# 将标签转换为Tensor
if len(boxes) > 0:
targets = torch.zeros((len(boxes), 6)) # (batch_index, class_id, xc, yc, w, h)
targets[:, 1] = torch.tensor(class_labels)
targets[:, 2:] = torch.tensor(boxes)
else:
targets = torch.zeros((0, 6))
return image, targets
这个 Dataset 类有几个设计考量：一是使用Albumentations处理YOLO格式的边界框非常方便；二是通过 min_visibility 参数，在增强（如裁剪）后自动过滤掉可见面积太小的目标，避免引入噪声；三是返回的 targets 张量第一列预留了 batch_index，这在后续计算损失时用于区分同一批次中不同图像的标签。
3. 模型架构深度实现：拆解YOLOV11的核心模块
网上很多文章对YOLOV11网络结构的描述停留在“骨干-颈部-头部”的层面，但真正实现时，细节决定成败。我们依据最新的公开资料和代码分析，来构建一个清晰、模块化的YOLOV11。
骨干网络（Backbone）：YOLOV11的骨干基于改进的CSPDarknet。其核心是CSP（Cross Stage Partial）结构，它通过将特征图通道拆分，一部分进行深度卷积处理，另一部分直接短路连接，有效减少了计算量并缓解了梯度消失。下面我们实现一个基础的CSPBlock：
import torch.nn as nn
class ConvBnAct(nn.Module):
"""标准卷积块：Conv2d + BatchNorm2d + SiLU激活"""
def __init__(self, in_c, out_c, kernel_size=1, stride=1, padding=None, groups=1):
super().__init__()
if padding is None:
padding = kernel_size // 2
self.conv = nn.Conv2d(in_c, out_c, kernel_size, stride, padding, groups=groups, bias=False)
self.bn = nn.BatchNorm2d(out_c)
self.act = nn.SiLU() # YOLOV11常用SiLU（Swish）激活函数
def forward(self, x):
return self.act(self.bn(self.conv(x)))
class Bottleneck(nn.Module):
"""标准瓶颈层，可选是否使用残差连接"""
def __init__(self, in_c, out_c, shortcut=True):
super().__init__()
hidden_c = out_c // 2
self.cv1 = ConvBnAct(in_c, hidden_c, 1)
self.cv2 = ConvBnAct(hidden_c, out_c, 3)
self.add = shortcut and in_c == out_c
def forward(self, x):
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
class CSPBlock(nn.Module):
"""CSP结构：特征图拆分->部分处理->拼接"""
def __init__(self, in_c, out_c, n_bottlenecks=1, shortcut=True):
super().__init__()
hidden_c = out_c // 2
# 主路径：1x1卷积降维后接多个Bottleneck
self.cv1 = ConvBnAct(in_c, hidden_c, 1)
self.m = nn.Sequential(*[Bottleneck(hidden_c, hidden_c, shortcut) for _ in range(n_bottlenecks)])
# 旁路：仅一个1x1卷积
self.cv2 = ConvBnAct(in_c, hidden_c, 1)
# 输出融合
self.cv3 = ConvBnAct(2 * hidden_c, out_c, 1)
def forward(self, x):
y1 = self.m(self.cv1(x))
y2 = self.cv2(x)
return self.cv3(torch.cat((y1, y2), dim=1))
有了CSPBlock，我们就可以搭建骨干网络了。YOLOV11的骨干通常由多个“阶段”（Stage）组成，每个阶段包含一个下采样卷积和一个或多个CSPBlock。
class Backbone(nn.Module):
def __init__(self, depth_multiple=1.0, width_multiple=1.0):
super().__init__()
# 基础通道数
base_channels = int(64 * width_multiple)
# 深度系数，控制每个Stage中CSPBlock的重复次数
depth_dict = {1: 1, 2: 2, 3: 8, 4: 8, 5: 4} # 示例配置，实际以官方为准
depth_dict = {k: max(round(v * depth_multiple), 1) for k, v in depth_dict.items()}
# Stem: 快速下采样
self.stem = nn.Sequential(
ConvBnAct(3, base_channels, 3, 2), # /2
ConvBnAct(base_channels, base_channels*2, 3, 2), # /4
)
# Stage 1-5
self.stages = nn.ModuleList()
in_c = base_channels * 2
for i in range(1, 6):
out_c = base_channels * (2 ** (i+1)) # 通道数逐级翻倍：128, 256, 512, 1024, 2048?
# 实际YOLOV11可能不会到2048，这里仅为示例
if i == 5:
out_c = int(out_c * 0.5) # 最后一层可能减少通道
stage = nn.Sequential(
ConvBnAct(in_c, out_c, 3, 2), # 下采样
CSPBlock(out_c, out_c, n_bottlenecks=depth_dict.get(i, 1))
)
self.stages.append(stage)
in_c = out_c
def forward(self, x):
features = []
x = self.stem(x)
features.append(x) # 输出1: 1/4尺度
for i, stage in enumerate(self.stages):
x = stage(x)
if i >= 1: # 通常取后三个Stage的输出给Neck
features.append(x)
return features # 返回多尺度特征图列表
颈部网络（Neck）：YOLOV11的颈部网络融合了SPPF（快速空间金字塔池化）和PANet（路径聚合网络）的思想，进行多尺度特征融合。SPPF模块是对SPP的优化，用串行的最大池化代替并行，减少计算量。
class SPPF(nn.Module):
"""快速空间金字塔池化"""
def __init__(self, in_c, out_c, kernel_sizes=5):
super().__init__()
hidden_c = in_c // 2
self.cv1 = ConvBnAct(in_c, hidden_c, 1)
# 多个串行的最大池化层
self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=k, stride=1, padding=k//2) for k in [kernel_sizes] * 3])
self.cv2 = ConvBnAct(hidden_c * 4, out_c, 1) # 拼接后通道数x4
def forward(self, x):
x = self.cv1(x)
y = [x]
for m in self.m:
x = m(x)
y.append(x)
return self.cv2(torch.cat(y, dim=1))
class PANet(nn.Module):
"""路径聚合网络：自顶向下 + 自底向上"""
def __init__(self, channels_list, depth_multiple=1.0):
super().__init__()
# 假设channels_list是骨干网络输出的三个特征图的通道数，例如[256, 512, 1024]
self.upsamples = nn.ModuleList()
self.downsamples = nn.ModuleList()
self.lateral_convs = nn.ModuleList()
# 构建自顶向下路径（上采样+融合）
for i in range(len(channels_list)-1):
in_c = channels_list[i+1]
out_c = channels_list[i]
self.upsamples.append(nn.Sequential(
ConvBnAct(in_c, out_c, 1),
nn.Upsample(scale_factor=2, mode='nearest')
))
self.lateral_convs.append(ConvBnAct(out_c*2, out_c, 3)) # 融合后卷积
# 构建自底向上路径（下采样+融合）
for i in range(len(channels_list)-1):
in_c = channels_list[i]
out_c = channels_list[i+1]
self.downsamples.append(nn.Sequential(
ConvBnAct(in_c, out_c, 3, 2), # stride=2实现下采样
))
# 此处的融合卷积可以与自顶向下路径共享，也可独立，这里独立定义
self.lateral_convs.append(ConvBnAct(out_c*2, out_c, 3))
def forward(self, features):
# features: [f1, f2, f3], 分辨率从大到小
# 自顶向下
laterals = [features[-1]]
for i in range(len(features)-2, -1, -1):
up_feat = self.upsamples[i](laterals[-1])
cat_feat = torch.cat([features[i], up_feat], dim=1)
laterals.append(self.lateral_convs[i](cat_feat))
laterals = laterals[::-1] # 反转，使其与输入顺序一致
# 自底向上
outputs = [laterals[0]]
for i in range(len(laterals)-1):
down_feat = self.downsamples[i](outputs[-1])
cat_feat = torch.cat([laterals[i+1], down_feat], dim=1)
outputs.append(self.lateral_convs[len(features)-1 + i](cat_feat))
return outputs # 返回融合后的多尺度特征
检测头（Head）：YOLOV11采用了解耦头（Decoupled Head），将分类和回归任务分开，这被证明能提升性能。同时，它依然是多尺度预测。
class YOLOV11Head(nn.Module):
def __init__(self, num_classes, in_channels_list, anchors_per_scale=3):
super().__init__()
self.num_classes = num_classes
self.nl = len(in_channels_list) # 尺度数量，通常是3
self.na = anchors_per_scale
# 为每个预测尺度创建独立的解耦头
self.cls_convs = nn.ModuleList() # 分类分支卷积
self.reg_convs = nn.ModuleList() # 回归分支卷积
self.cls_preds = nn.ModuleList() # 分类预测层
self.reg_preds = nn.ModuleList() # 回归预测层
for in_c in in_channels_list:
# 分类分支：两个卷积后接预测层
cls_conv = nn.Sequential(
ConvBnAct(in_c, in_c, 3),
ConvBnAct(in_c, in_c, 3),
)
self.cls_convs.append(cls_conv)
self.cls_preds.append(nn.Conv2d(in_c, self.na * self.num_classes, 1))
# 回归分支：两个卷积后接预测层（预测框坐标+objectness）
reg_conv = nn.Sequential(
ConvBnAct(in_c, in_c, 3),
ConvBnAct(in_c, in_c, 3),
)
self.reg_convs.append(reg_conv)
# 输出：每个anchor预测4个坐标偏移量(tx, ty, tw, th) + 1个objectness置信度
self.reg_preds.append(nn.Conv2d(in_c, self.na * 5, 1))
def forward(self, features):
"""features: 来自Neck的多尺度特征图列表"""
cls_outputs, reg_outputs = [], []
for i, feat in enumerate(features):
cls_feat = self.cls_convs[i](feat)
reg_feat = self.reg_convs[i](feat)
cls_out = self.cls_preds[i](cls_feat) # [B, na*C, H, W]
reg_out = self.reg_preds[i](reg_feat) # [B, na*5, H, W]
# 调整形状为 [B, na, H, W, ...] 以便后续处理
B, _, H, W = cls_out.shape
cls_out = cls_out.view(B, self.na, self.num_classes, H, W).permute(0, 1, 3, 4, 2).contiguous()
reg_out = reg_out.view(B, self.na, 5, H, W).permute(0, 1, 3, 4, 2).contiguous()
cls_outputs.append(cls_out)
reg_outputs.append(reg_out)
return torch.cat([v.flatten(1, 3) for v in cls_outputs], dim=1), \
torch.cat([v.flatten(1, 3) for v in reg_outputs], dim=1)
最后，我们将所有模块组装成完整的YOLOV11模型：
class YOLOV11(nn.Module):
def __init__(self, num_classes=80, depth_multiple=1.0, width_multiple=1.0):
super().__init__()
self.backbone = Backbone(depth_multiple, width_multiple)
# 假设Backbone输出三个尺度的特征图，通道数为[256, 512, 1024]（根据width_multiple缩放）
backbone_out_channels = [int(256*width_multiple), int(512*width_multiple), int(1024*width_multiple)]
# Neck: 先SPPF，再PANet
self.sppf = SPPF(backbone_out_channels[-1], backbone_out_channels[-1])
# 更新最后一个特征图的通道（SPPF可能改变通道数）
neck_in_channels = backbone_out_channels.copy()
neck_in_channels[-1] = backbone_out_channels[-1] # SPPF输出通道不变
self.neck = PANet(neck_in_channels, depth_multiple)
# PANet输出通道数与输入一致
self.head = YOLOV11Head(num_classes, neck_in_channels)
def forward(self, x):
backbone_feats = self.backbone(x) # 多尺度特征
# 对最后一个特征图应用SPPF
backbone_feats[-1] = self.sppf(backbone_feats[-1])
neck_feats = self.neck(backbone_feats)
cls_pred, reg_pred = self.head(neck_feats)
return cls_pred, reg_pred # 返回分类和回归预测
4. 训练策略与损失函数：让模型真正学会“检测”
模型搭好了，但如何训练它？YOLO的损失函数是其灵魂，它需要同时处理分类、边界框回归和对象置信度。YOLOV11的损失函数通常包含三个部分：分类损失、边界框回归损失和目标置信度损失。其中，边界框回归损失近年来从简单的L1/L2损失进化到了IoU系列损失，如CIoU、DIoU，它们能更好地衡量框的重叠度和中心点距离。
我们先实现一个计算IoU和CIoU的函数：
def bbox_iou(box1, box2, xywh=True, eps=1e-7):
"""计算两个边界框集合之间的IoU (Intersection over Union)"""
if xywh:
# 将 (x_center, y_center, width, height) 转换为 (x1, y1, x2, y2)
b1_x1, b1_x2 = box1[..., 0] - box1[..., 2] / 2, box1[..., 0] + box1[..., 2] / 2
b1_y1, b1_y2 = box1[..., 1] - box1[..., 3] / 2, box1[..., 1] + box1[..., 3] / 2
b2_x1, b2_x2 = box2[..., 0] - box2[..., 2] / 2, box2[..., 0] + box2[..., 2] / 2
b2_y1, b2_y2 = box2[..., 1] - box2[..., 3] / 2, box2[..., 1] + box2[..., 3] / 2
else:
b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, dim=-1)
b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, dim=-1)
b1_x1, b1_y1, b1_x2, b1_y2 = b1_x1.squeeze(-1), b1_y1.squeeze(-1), b1_x2.squeeze(-1), b1_y2.squeeze(-1)
b2_x1, b2_y1, b2_x2, b2_y2 = b2_x1.squeeze(-1), b2_y1.squeeze(-1), b2_x2.squeeze(-1), b2_y2.squeeze(-1)
# 交集区域
inter_x1 = torch.max(b1_x1, b2_x1)
inter_y1 = torch.max(b1_y1, b2_y1)
inter_x2 = torch.min(b1_x2, b2_x2)
inter_y2 = torch.min(b1_y2, b2_y2)
inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * torch.clamp(inter_y2 - inter_y1, min=0)
# 并集区域
b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
union_area = b1_area + b2_area - inter_area + eps
iou = inter_area / union_area
return iou
def bbox_ciou(box1, box2, eps=1e-7):
"""计算CIoU (Complete IoU) Loss"""
iou = bbox_iou(box1, box2, xywh=True, eps=eps)
# 中心点距离的平方
cw = torch.max(box1[..., 0], box2[..., 0]) - torch.min(box1[..., 0], box2[..., 0])
ch = torch.max(box1[..., 1], box2[..., 1]) - torch.min(box1[..., 1], box2[..., 1])
c2 = cw ** 2 + ch ** 2 + eps # 对角线距离平方
rho2 = ((box1[..., 0] - box2[..., 0]) ** 2 + (box1[..., 1] - box2[..., 1]) ** 2) # 中心点距离平方
# 宽高比一致性
w1, h1 = box1[..., 2], box1[..., 3]
w2, h2 = box2[..., 2], box2[..., 3]
v = (4 / (torch.pi ** 2)) * torch.pow(torch.atan(w2 / (h2 + eps)) - torch.atan(w1 / (h1 + eps)), 2)
with torch.no_grad():
alpha = v / (v - iou + (1 + eps))
ciou = iou - (rho2 / c2 + v * alpha)
return ciou
基于此，我们可以构建YOLO的损失函数。这里需要注意正负样本的分配策略。YOLO通常采用基于Anchor和GT（真实框）IoU的分配方式，但最新的方法可能使用Task-Aligned Assigner等更复杂的策略。为了简化，我们实现一个基于IoU阈值的版本：
class YOLOLoss(nn.Module):
def __init__(self, num_classes, anchors, img_size=640):
super().__init__()
self.num_classes = num_classes
self.anchors = anchors # 形状为 [nl, na, 2]，每个尺度每个anchor的宽高
self.img_size = img_size
self.bce_cls = nn.BCEWithLogitsLoss(reduction='none')
self.bce_obj = nn.BCEWithLogitsLoss(reduction='none')
# 损失分量权重
self.lambda_cls = 0.5
self.lambda_box = 0.05
self.lambda_obj = 1.0
def forward(self, predictions, targets):
"""
predictions: 元组 (cls_pred, reg_pred)
cls_pred: [B, N, num_classes]
reg_pred: [B, N, 5] (tx, ty, tw, th, obj)
targets: [B, M, 6] (image_idx, class_id, xc, yc, w, h)
"""
cls_pred, reg_pred = predictions
B, N, _ = cls_pred.shape
# 初始化损失
loss_cls = torch.zeros(1, device=cls_pred.device)
loss_box = torch.zeros(1, device=cls_pred.device)
loss_obj = torch.zeros(1, device=cls_pred.device)
# 将预测解码为边界框 (需要根据grid和anchor进行解码，此处简化流程)
# 实际实现中，需要根据特征图网格位置、anchor大小等解码出预测框的绝对坐标
# 这里省略了解码和分配的具体代码，它是一个相对复杂的过程
# 伪代码：正负样本分配
# 1. 为每个GT框找到IoU最大的Anchor
# 2. 将对应位置标记为正样本
# 3. 计算分类损失（仅正样本）
# 4. 计算边界框CIoU损失（仅正样本）
# 5. 计算目标置信度损失（正样本为1，负样本为0，使用二元交叉熵）
# 假设我们已经得到了正样本掩码 pos_mask [B, N] 和匹配的GT框 matched_gt [B, N, 4]
# 以及分类标签 cls_target [B, N, num_classes]
# 和对象置信度标签 obj_target [B, N]
# 分类损失
if pos_mask.sum() > 0:
loss_cls = self.bce_cls(cls_pred[pos_mask], cls_target[pos_mask]).mean()
# 边界框回归损失
if pos_mask.sum() > 0:
# 解码预测框
pred_boxes = self.decode_boxes(reg_pred[..., :4]) # 需要实现decode_boxes
# 计算CIoU
ciou = bbox_ciou(pred_boxes[pos_mask], matched_gt[pos_mask])
loss_box = (1.0 - ciou).mean()
# 对象置信度损失
loss_obj_pos = self.bce_obj(reg_pred[..., 4][pos_mask], obj_target[pos_mask]).mean() if pos_mask.sum() > 0 else 0
loss_obj_neg = self.bce_obj(reg_pred[..., 4][~pos_mask], obj_target[~pos_mask]).mean() if (~pos_mask).sum() > 0 else 0
loss_obj = loss_obj_pos + 0.5 * loss_obj_neg # 负样本权重可调
total_loss = self.lambda_cls * loss_cls + self.lambda_box * loss_box + self.lambda_obj * loss_obj
return total_loss, {'cls': loss_cls.item(), 'box': loss_box.item(), 'obj': loss_obj.item()}
有了损失函数，我们就可以开始训练循环了。训练中的一些关键技巧包括：
- 优化器选择：AdamW现在比SGD更受欢迎，尤其是配合余弦退火学习率调度。
- 热身（Warmup）：在训练初期使用较低的学习率，有助于稳定训练。
- 梯度累积：当GPU内存不足时，可以通过多次前向传播累积梯度，再一次性更新参数，模拟更大的批次大小。
- 混合精度训练：使用
torch.cuda.amp可以显著减少显存占用并加快训练速度。
下面是一个训练循环的核心框架：
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR
from torch.cuda.amp import GradScaler, autocast
def train_one_epoch(model, train_loader, criterion, optimizer, scheduler, scaler, epoch, device):
model.train()
total_loss = 0
pbar = tqdm(train_loader, desc=f'Epoch {epoch}')
for batch_idx, (images, targets) in enumerate(pbar):
images = images.to(device)
targets = targets.to(device)
# 混合精度训练上下文
with autocast():
predictions = model(images)
loss, loss_dict = criterion(predictions, targets)
# 梯度缩放与反向传播
scaler.scale(loss).backward()
# 梯度累积：每4个批次更新一次参数
if (batch_idx + 1) % 4 == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
total_loss += loss.item()
pbar.set_postfix({'loss': loss.item(), 'cls': loss_dict['cls'], 'box': loss_dict['box'], 'obj': loss_dict['obj']})
scheduler.step() # 更新学习率
return total_loss / len(train_loader)
# 主训练函数
def main():
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = YOLOV11(num_classes=80).to(device)
# 数据加载
train_dataset = YOLODataset('data/coco/train2017', 'data/coco/labels/train2017', augment=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=4, collate_fn=collate_fn)
# 优化器与调度器
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=5e-4)
# 线性热身
warmup_scheduler = LinearLR(optimizer, start_factor=0.1, total_iters=3)
# 余弦退火主调度
main_scheduler = CosineAnnealingLR(optimizer, T_max=100-3, eta_min=1e-5)
# 组合调度器
from torch.optim.lr_scheduler import SequentialLR
scheduler = SequentialLR(optimizer, schedulers=[warmup_scheduler, main_scheduler], milestones=[3])
criterion = YOLOLoss(num_classes=80, anchors=[[10,13], [16,30], [33,23], [30,61], [62,45], [59,119], [116,90], [156,198], [373,326]])
scaler = GradScaler()
for epoch in range(100):
avg_loss = train_one_epoch(model, train_loader, criterion, optimizer, scheduler, scaler, epoch, device)
print(f'Epoch {epoch} Average Loss: {avg_loss:.4f}')
# 每隔一定epoch保存检查点
if (epoch + 1) % 10 == 0:
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': avg_loss,
}, f'checkpoints/yolov11_epoch_{epoch}.pth')
5. 模型推理与后处理：从输出到可视化结果
训练好的模型最终要用于推理。推理过程包括前向传播、解码预测框、应用非极大值抑制（NMS）过滤冗余框，最后将结果可视化。
首先，我们需要将模型输出的偏移量解码为图像上的绝对坐标。这需要知道每个预测对应的网格位置和预设的Anchor大小。
def decode_predictions(predictions, anchors, img_size, num_classes, conf_thresh=0.25):
"""
将模型原始输出解码为可读的边界框。
predictions: (cls_pred, reg_pred)
anchors: 预设的Anchor框，形状为 [nl, na, 2]
返回: list of [x1, y1, x2, y2, confidence, class_id] for each image in batch
"""
cls_pred, reg_pred = predictions
B = cls_pred.shape[0]
decoded_results = []
# 这里需要根据模型输出的特征图网格大小，生成网格坐标
# 并利用anchors将预测的偏移量(tx, ty, tw, th)转换为绝对坐标(xc, yc, w, h)
# 再将中心坐标格式转换为角点坐标格式(x1, y1, x2, y2)
# 同时应用sigmoid到分类置信度和对象置信度，并过滤掉低置信度的预测
# 这是一个简化的解码流程示意，实际实现需要考虑多尺度
for i in range(B):
boxes = []
scores = []
class_ids = []
# 遍历每个尺度的预测...
# 假设我们得到了一个初步的预测列表 all_boxes [N, 6]
# 其中每行: [x1, y1, x2, y2, obj_conf, class_conf, class_id]
# 应用置信度阈值
keep = all_boxes[:, 4] * all_boxes[:, 5] > conf_thresh # obj_conf * class_conf
filtered_boxes = all_boxes[keep]
if len(filtered_boxes) == 0:
decoded_results.append(torch.zeros((0, 6), device=all_boxes.device))
continue
# 非极大值抑制 (NMS)
nms_indices = nms(filtered_boxes[:, :4], filtered_boxes[:, 4] * filtered_boxes[:, 5], iou_threshold=0.45)
final_boxes = filtered_boxes[nms_indices]
decoded_results.append(final_boxes)
return decoded_results
def nms(boxes, scores, iou_threshold):
"""非极大值抑制"""
# boxes: [N, 4] (x1, y1, x2, y2)
# scores: [N]
if boxes.numel() == 0:
return torch.empty((0,), dtype=torch.long, device=boxes.device)
x1 = boxes[:, 0]
y1 = boxes[:, 1]
x2 = boxes[:, 2]
y2 = boxes[:, 3]
areas = (x2 - x1) * (y2 - y1)
# 按分数降序排序
_, order = scores.sort(descending=True)
keep = []
while order.numel() > 0:
i = order[0]
keep.append(i)
if order.numel() == 1:
break
# 计算当前框与剩余框的IoU
xx1 = torch.max(x1[i], x1[order[1:]])
yy1 = torch.max(y1[i], y1[order[1:]])
xx2 = torch.min(x2[i], x2[order[1:]])
yy2 = torch.min(y2[i], y2[order[1:]])
w = torch.clamp(xx2 - xx1, min=0)
h = torch.clamp(yy2 - yy1, min=0)
inter = w * h
iou = inter / (areas[i] + areas[order[1:]] - inter)
# 保留IoU低于阈值的框
idx = torch.where(iou <= iou_threshold)[0]
order = order[idx + 1] # +1 因为order[1:]比order少一个元素
return torch.tensor(keep, device=boxes.device)
最后，我们可以编写一个推理脚本，加载训练好的权重，对单张图像或整个文件夹进行预测并可视化结果。
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
def predict_and_visualize(model, img_path, conf_thresh=0.25, iou_thresh=0.45):
model.eval()
device = next(model.parameters()).device
# 加载并预处理图像
orig_img = cv2.imread(img_path)
img_rgb = cv2.cvtColor(orig_img, cv2.COLOR_BGR2RGB)
H, W = img_rgb.shape[:2]
# 预处理：resize, normalize, to tensor
img_tensor = preprocess_image(img_rgb, img_size=640).to(device).unsqueeze(0)
with torch.no_grad():
predictions = model(img_tensor)
detections = decode_predictions(predictions, anchors, img_size=640, num_classes=80, conf_thresh=conf_thresh)
# 将检测框绘制到原图上
img_pil = Image.fromarray(img_rgb)
draw = ImageDraw.Draw(img_pil)
# 可以加载一个字体文件使标签更美观
# font = ImageFont.truetype("arial.ttf", 16)
for det in detections[0]: # 批次中第一张图
x1, y1, x2, y2, conf, cls_id = det.cpu().numpy()
# 将归一化坐标转换回原图尺寸
x1, x2 = int(x1 * W), int(x2 * W)
y1, y2 = int(y1 * H), int(y2 * H)
# 绘制矩形框
draw.rectangle([x1, y1, x2, y2], outline='red', width=2)
# 绘制标签背景和文字
label = f"Class {int(cls_id)}: {conf:.2f}"
text_bbox = draw.textbbox((x1, y1-20), label) # 简单估算文字区域
draw.rectangle(text_bbox, fill='red')
draw.text((x1, y1-20), label, fill='white') #, font=font)
# 显示或保存结果
result_img = np.array(img_pil)
cv2.imshow('Detection Result', cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR))
cv2.waitKey(0)
cv2.destroyAllWindows()
# 或者保存
# cv2.imwrite('result.jpg', cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR))
从环境搭建到模型推理，我们走完了YOLOV11项目实战的完整闭环。这个过程涉及大量细节，从数据管道的高效构建、模型模块的灵活设计，到损失函数中正负样本分配的微妙平衡，每一步都需要仔细推敲。我提供的代码是一个高度模块化的起点，你可以在此基础上进行修改和实验，例如尝试不同的数据增强组合、调整网络深度和宽度、或者集成更先进的损失函数如Focal Loss。在实际项目中，你很可能需要根据自己数据集的特性（如目标尺度分布、类别不平衡问题）来调整Anchor大小、数据增强策略甚至网络结构。记住，读懂代码并动手调试，远比复制粘贴更能让你理解模型是如何工作的。
DAMO开发者矩阵，由阿里巴巴达摩院和中国互联网协会联合发起，致力于探讨最前沿的技术趋势与应用成果，搭建高质量的交流与分享平台，推动技术创新与产业应用链接，围绕“人工智能与新型计算”构建开放共享的开发者生态。
更多推荐
2
0- 0
已为社区贡献5条内容
所有评论(0)