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注意力机制 

来源地址：https://blog.csdn.net/qq_49627063/article/details/130015203