071、YOLOv8改进实战旋转框检测头设计与角度回归损失函数KLD/SkewIoU的代码集成从一次无人机航拍项目说起去年接了个无人机检测项目甲方要求检测任意朝向的车辆和集装箱。一开始用标准YOLOv8跑水平框检测效果还行但遇到斜着停的车、堆场里歪着的集装箱mAP直接掉到0.3以下。更头疼的是水平框把两个挨着的斜目标框在一起NMS后漏检严重。那段时间我盯着可视化结果满屏的红色框框得乱七八糟心里只有一个念头必须上旋转框。如果你也遇到过类似场景——遥感图像、工业零件检测、文档版面分析、任意朝向的文本检测——那么YOLOv8加旋转框检测头这条路你迟早得走一遍。今天这篇笔记我就把当时踩过的坑、试过的方案、最终落地的代码原原本本写出来。旋转框检测头从0到1的改造YOLOv8的检测头输出是(batch, 4num_classes, grid_h, grid_w)4代表(x, y, w, h)。旋转框需要多一个角度参数变成5个回归量(x, y, w, h, theta)。theta的范围我见过用[-π/4, 3π/4)的也见过[-π/2, π/2)的还有用[0, π)的。这里我踩过坑——不同范围直接影响损失函数的收敛行为后面讲KLD的时候细说。先看检测头改造的核心代码。YOLOv8的检测头在ultralytics/nn/modules/head.py里Detect类。我直接复制了一份叫RotateDetect改动点就两个输出通道数从4num_classes改成5num_classes以及解码时把角度加上。classRotateDetect(nn.Module):def__init__(self,nc80,ch()):super().__init__()self.ncnc# 类别数self.nonc5# 每个anchor的输出x,y,w,h,theta classself.nllen(ch)# 检测层数一般是3层self.stridetorch.zeros(self.nl)# 后面会赋值# 每个检测层的卷积输出通道是 self.no * self.reg_max# 注意这里 reg_max 我设成了16和原版保持一致self.cv2nn.ModuleList(nn.Sequential(Conv(x,x,3),Conv(x,x,3),nn.Conv2d(x,self.no*16,1)# 别这样写直接乘16后面reshape)forxinch)# 分类分支不变self.cv3nn.ModuleList(nn.Sequential(Conv(x,x,3),Conv(x,x,3),nn.Conv2d(x,self.nc*16,1))forxinch)这里有个细节原版YOLOv8的回归分支输出是4 * reg_max用DFLDistribution Focal Loss做分布回归。我保留了DFL结构把4改成5角度也做分布回归。但后来实验发现角度用DFL效果并不好——角度是周期性的DFL的softmax分布没法很好地表达周期性。最终我改成了直接回归一个值不经过DFL。踩坑记录如果你保留DFL给角度训练时角度预测值会在边界处震荡比如真实角度是179度假设范围0-180模型预测在0度附近损失巨大。后来我改成角度分支单独用Smooth L1 Loss问题解决。解码逻辑把网络输出转成旋转框解码部分原版YOLOv8用dist2bbox把分布转成边界框。我写了个dist2rbbox把角度单独拎出来。defdist2rbbox(distance,anchor_points,xywhTrue,dim1):将DFL输出转换为旋转框坐标# distance shape: (batch, 5*16, h, w) - 5组分布每组16个binbatch,_,h,wdistance.shape distancedistance.view(batch,5,16,h,w)# 对前4个通道做DFL解码x, y, w, hprojtorch.arange(16,dtypedistance.dtype,devicedistance.device)dflF.softmax(distance[:,:4],dim2) proj# (batch, 4, h, w)# 角度直接取分布加权平均或者用argmax# 这里踩过坑用argmax梯度传不过去训练时角度不更新# 改用softmax加权平均虽然有点近似但梯度能传angle_probsF.softmax(distance[:,4:5],dim2)# (batch, 1, 16, h, w)theta(angle_probs proj).squeeze(1)# (batch, h, w)# 归一化到 [-pi/4, 3pi/4)thetatheta/16*np.pi-np.pi/4# 映射到目标范围# 后面跟原版一样用anchor_points解码x,y,w,h# ...别这样写把theta直接sigmoid然后乘范围。我试过角度分布不均匀靠近边界的地方预测不准。用DFL的分布形式让网络自己学角度的概率分布效果更好。角度回归损失函数KLD和SkewIoU的恩怨旋转框的损失函数我前后试了三种Smooth L1 on theta、KLD、SkewIoU。先说结论KLD最稳SkewIoU精度上限高但训练不稳定Smooth L1适合快速验证。KLDKullback-Leibler Divergence损失KLD的核心思想是把旋转框看成二维高斯分布用KL散度衡量两个分布的差异。这样做的好处是角度、位置、尺寸的误差被统一到同一个度量空间不需要手动调权重。defkld_loss(pred_rbboxes,target_rbboxes,reductionmean): pred_rbboxes: (N, 5) 格式 [x, y, w, h, theta] target_rbboxes: (N, 5) theta范围: [-pi/4, 3pi/4) # 把旋转框转成高斯分布的均值向量和协方差矩阵# 这里踩过坑直接用公式算注意theta是弧度pred_gaussiansrbbox_to_gaussian(pred_rbboxes)target_gaussiansrbbox_to_gaussian(target_rbboxes)# 计算KL散度mu1,sigma1pred_gaussians mu2,sigma2target_gaussians# 公式KL(N1||N2) 0.5 * (tr(sigma2^-1 sigma1) (mu2-mu1)^T sigma2^-1 (mu2-mu1) - 2 ln(det(sigma2)/det(sigma1)))# 直接实现注意数值稳定性sigma2_invtorch.inverse(sigma2)term1torch.bmm(sigma2_inv,sigma1).diagonal(dim1-2,dim2-1).sum(-1)diffmu2-mu1 term2torch.bmm(diff.unsqueeze(1),torch.bmm(sigma2_inv,diff.unsqueeze(-1))).squeeze()term3torch.log(torch.det(sigma2)/torch.det(sigma1).clamp(min1e-8))kl0.5*(term1term2-2term3)returnkl.mean()ifreductionmeanelseklKLD有个坑当预测框和真实框差异很大时KL散度可能为负理论上KL散度非负但数值计算可能出负值。我加了个clamp(min0)别直接取绝对值否则梯度方向会乱。SkewIoU损失SkewIoU就是旋转框的IoU但计算起来比水平框复杂得多需要算两个旋转矩形的交集面积。我用了mmrotate里的实现但发现一个问题SkewIoU在角度误差大时梯度不稳定尤其是两个框完全不重叠时IoU0梯度为0模型学不动。defskew_iou_loss(pred_rbboxes,target_rbboxes): 用mmrotate的旋转框IoU计算 注意这个函数在batch size大时很慢因为要逐框计算 frommmrotate.coreimportobb_overlaps iousobb_overlaps(pred_rbboxes.detach().cpu().numpy(),target_rbboxes.detach().cpu().numpy(),iou)ioustorch.from_numpy(ious).to(pred_rbboxes.device)# 这里踩过坑直接返回1-iou作为损失梯度回传有问题# 因为obb_overlaps是numpy操作梯度断掉了# 解决方案用可微的近似IoU或者把SkewIoU作为辅助损失return(1-ious).mean()别这样写直接调mmrotate的numpy版本梯度断了。正确做法是用可微的SkewIoU实现或者像我一样把SkewIoU作为评估指标训练时用KLD。训练技巧让旋转框收敛的秘诀旋转框训练比水平框敏感得多我总结几个关键点学习率要降。同样的模型水平框用0.01能收敛旋转框可能炸掉。我降到0.001起步用余弦退火调度。角度初始化。检测头最后的卷积层角度分支的bias初始化为0对应theta0水平框。这样模型一开始先学水平检测再慢慢学旋转。直接随机初始化前期损失巨大。数据增强要小心。Mosaic和MixUp会生成奇怪的旋转框组合尤其是Mosaic把四个图的旋转框拼在一起角度分布变得混乱。我最后只保留了随机旋转-30度到30度和随机裁剪其他增强都关了。NMS后处理。旋转框的NMS不能用普通IoU得用SkewIoU。我写了个rotate_nms用torchvision.ops.nms的接口但IoU计算换成旋转框版本。注意阈值要调低我一般设0.3比水平框的0.5低。代码集成把旋转框塞进YOLOv8训练流程改完检测头和损失函数还要改数据加载、评估指标、可视化。这里说几个容易漏的点数据加载。YOLO格式的标注是class x_center y_center width height旋转框要改成class x_center y_center width height theta。注意theta的单位和范围要统一。我用的DOTA数据集格式theta是角度制范围0-180但模型里用弧度所以加载时要转换。评估指标。mAP计算要用旋转框的IoU。我直接用了mmrotate里的评估代码但注意它要求输入是(x, y, w, h, theta)格式theta范围[-π/2, π/2)。如果模型输出范围不同要转换。可视化。OpenCV的rectangle只能画水平框。我写了个函数用cv2.boxPoints把旋转框转成四个角点然后用polylines画。别直接用cv2.minAreaRect它返回的是中心点尺寸角度但角度范围是[-90, 0)跟模型输出不一致。经验性建议如果你现在要改YOLOv8做旋转框检测我的建议是先跑通KLD版本别一上来就上SkewIoU。KLD训练稳定收敛快虽然精度上限可能比SkewIoU低1-2个点但至少能让你看到模型在学。等KLD版本跑通了再尝试用SkewIoU做fine-tune或者两个损失加权组合。角度范围选[-π/4, 3π/4)这个范围的好处是水平框对应0度垂直框对应π/2长条形的目标不会因为角度周期性导致歧义。[-π/2, π/2)也行但要注意边界处理。最后别指望旋转框能解决所有问题。如果你的目标长宽比接近1:1比如正方形旋转框带来的提升有限反而增加计算量。旋转框最适合长条形、任意朝向的目标——车辆、船舶、集装箱、文本行。那次无人机项目最终用KLD损失旋转框检测头mAP从0.3提到了0.62。虽然离甲方要求的0.7还有差距但至少从不能用变成了勉强能用。后来加了多尺度训练和测试时增强才勉强达标。旋转框这条路坑多但值得走。
