yoloV5改进-针对小目标的NWD

 论文地址：https://arxiv.org/pdf/2110.13389.pdf

目录

前言:

wasserstein_loss 

LOSS.PY

---

前言:

        最近在做目标检测时，图片分辨率为6016 x 2048，遇到一个200*200的小目标问题，这里加入了NWP针对该问题做出优化

wasserstein_loss 

utils/metrics.py中加入wasserstein_loss

def wasserstein_loss(pred, target, eps=1e-7, constant=12.8):
    """Implementation of paper `A Normalized Gaussian Wasserstein Distance for
    Tiny Object Detection .
    Args:
        pred (Tensor): Predicted bboxes of format (cx, cy, w, h),
            shape (n, 4).
        target (Tensor): Corresponding gt bboxes, shape (n, 4).
        eps (float): Eps to avoid log(0).
    Return:
        Tensor: Loss tensor.
    """
    center1 = pred[:, :2]
    center2 = target[:, :2]
    whs = center1[:, :2] - center2[:, :2]
    center_distance = whs[:, 0] * whs[:, 0] + whs[:, 1] * whs[:, 1] + eps

    w1 = pred[:, 2] + eps
    h1 = pred[:, 3] + eps
    w2 = target[:, 2] + eps
    h2 = target[:, 3] + eps

    wh_distance = ((w1 - w2) ** 2 + (h1 - h2) ** 2) / 4
    wasserstein_2 = center_distance + wh_distance
    return torch.exp(-torch.sqrt(wasserstein_2) / constant)

LOSS.PY

修改loss.py

            if n:
                # pxy, pwh, _, pcls = pi[b, a, gj, gi].tensor_split((2, 4, 5), dim=1)  # faster, requires torch 1.8.0
                pxy, pwh, _, pcls = pi[b, a, gj, gi].split((2, 2, 1, self.nc), 1)  # target-subset of predictions

                # Regression
                pxy = pxy.sigmoid() * 2 - 0.5
                pwh = (pwh.sigmoid() * 2) ** 2 * anchors[i]
                pbox = torch.cat((pxy, pwh), 1)  # predicted box
                iou = bbox_iou(pbox, tbox[i], CIoU=True).squeeze()  # iou(prediction, target)
                # ==================================================
                # lbox += (1.0 - iou).mean()  # iou loss 修改
                # Objectness
                # score_iou = iou.detach().clamp(0).type(tobj.dtype)  #修改

                nwd = wasserstein_loss(pbox, tbox[i]).squeeze()
                nwd_ratio = 0.5  # 平衡稀疏  nwd和 iou各取0.5   如果数据集全是小目标的换可以设置1或者0.9，,08 意思为只用nwd
                lbox += (1 - nwd_ratio) * (1.0 - nwd).mean() + nwd_ratio * (1.0 - iou).mean()
                # Objectness
                iou = (iou.detach() * nwd_ratio + nwd.detach() * (1 - nwd_ratio)).clamp(0, 1).type(tobj.dtype) #这里clamp（0，1）中设置了最小参数必须大于等于0进行
                # ===============================================

                # Objectness
                iou = iou.detach().clamp(0).type(tobj.dtype)

结语：

后续更新对小目标改进：有效BiFormer注意力机制