【YOLO改进】换遍IoU损失函数之SIoU Loss（基于MMYOLO）

五山一胖

已于 2024-04-30 07:54:05 修改

阅读量1.1k

点赞数 8

CC 4.0 BY-SA版权

分类专栏： YOLO改进文章标签： YOLO

于 2024-04-29 09:16:55 首次发布

本文链接：https://blog.youkuaiyun.com/Vlone_pp/article/details/138301125

SIoU损失函数

论文链接:https://arxiv.org/pdf/2205.12740

SIoU（Simplified IoU）损失函数是一种基于IoU（Intersection over Union）的改进损失函数，主要用于目标检测任务中的边界框回归。与传统的IoU损失函数相比，SIoU损失函数考虑了更多的几何特性，如边界框的中心点位置、长宽比和角度，从而提高了检测的准确性和鲁棒性。

设计思路

中心点距离：传统IoU损失主要关注边界框的重叠区域，但没有直接考虑框的中心位置。SIoU损失引入了一个中心点距离项，即预测框与真实框中心点之间的距离。这有助于优化模型以使预测框更准确地定位到目标的中心。
长宽比和角度：除了中心点距离外，SIoU还考虑了框的长宽比和角度的一致性。这是因为在很多实际应用中，目标的形状和方向信息对于正确的识别同样重要。通过对这些几何属性的考量，SIoU能够更好地处理不同形状和方向的目标。
简化计算：尽管引入了更多的几何特性，SIoU损失函数的设计还是尽可能地简化了计算过程，以便于在实际的深度学习框架中高效实现。

计算方式

计算中心点距离：
- 设定预测框中心为 $(x_p,y_p)$ ，真实框中心为 $(x_g,y_g)$ 。
- 中心点距离 D 可以使用欧氏距离计算： $D = \sqrt{(x_p - x_g)^2 + (y_p - y_g)^2}$
长宽比和角度误差：
- 计算预测框和真实框的长宽比 $r_p$ 和 $r_g$ 。
- 通过一定的函数（如对数差）计算长宽比误差。
- 若考虑角度，还需要计算框的旋转角度差。
结合不同项的权重：
- 对上述计算出的中心点距离、长宽比误差和可能的角度误差按照一定的权重进行加权，以形成最终的SIoU损失。权重可以根据具体任务和数据集的特性进行调整。
优化目标：
最小化SIoU损失，即通过梯度下降等优化算法调整模型参数，以使得预测框与真实框之间的SIoU损失最小化。通过这种方式，SIoU损失能够综合考虑目标的位置、形状和方向等多个方面，使得目标检测模型能够更全面地学习到目标的特性，从而提高检测性能。

基于Pytorch实现SIoU Loss

import torch
import torch.nn as nn

class SIoULoss(nn.Module):
    def __init__(self):
        super(SIoULoss, self).__init__()

    def forward(self, preds, targets):
        # preds和targets的形状应该是[N, 4]，其中每个框表示为[x_center, y_center, width, height]
        
        # 提取预测和真实的中心点坐标
        pred_centers = preds[:, :2]
        target_centers = targets[:, :2]
        
        # 提取预测和真实的宽度和高度
        pred_sizes = preds[:, 2:]
        target_sizes = targets[:, 2:]
        
        # 计算中心点的欧式距离
        center_distance = torch.norm(pred_centers - target_centers, dim=1)
        
        # 计算宽高比的对数损失
        pred_log_wh = torch.log(pred_sizes)
        target_log_wh = torch.log(target_sizes)
        wh_loss = torch.abs(pred_l