问题:

使用YOLOv5进行测试的时候,报错:Could not run 'torchvision::nms' with arguments from the 'CUDA' backend.(如下图所示)

 解决方法:

(1)按照网上绝大多数的做法,重新安装torch和torchvision,我的另外一篇博客有讲解,注意CUDA、torch、torchvision和其他可能的安装包的版本要对应,链接直达:https://blog.csdn.net/qq_54185421/article/details/124759003?spm=1001.2014.3001.5501https://blog.csdn.net/qq_54185421/article/details/124759003?spm=1001.2014.3001.5501(2)根据后台提示的信息,分析源码,进行问题的解决,下面是解决问题的过程:

(a) ctrl + 鼠标左键,点击:

(b)接下来看到如下界面:

 (c)分析代码行:return torch.ops.torchvision.nms(boxes, scores, iou_threshold),ctrl + 鼠标左键点击 nms,进入到下一个文件,截图如下图所示:

 (d)在函数nms内部做一个标志(print(22222222222222)),观察能否调用函数nms,经测试发现没有调用函数nms。

(e)总结:非极大值抑制这一环节出现了问题,导致检测的时候出现了错误。

(f)从官网或者其他网站找一个nms的代码(或者自己编写),插入到YOLOv5项目中。

(g)直到代码后,单独写成一个py文件(代码放到了文章的最后),取名my_nms.py,然后在bbox.py文件中调用,如下图所示:

 (h)运行结果图:

my_nms.py

from torch import Tensor
import torch


def box_area(boxes: Tensor) -> Tensor:
    """
    Computes the area of a set of bounding boxes, which are specified by its
    (x1, y1, x2, y2) coordinates.
    Arguments:
        boxes (Tensor[N, 4]): boxes for which the area will be computed. They
            are expected to be in (x1, y1, x2, y2) format
    Returns:
        area (Tensor[N]): area for each box
    """
    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])


def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
    """
    Return intersection-over-union (Jaccard index) of boxes.
    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
    Arguments:
        boxes1 (Tensor[N, 4])
        boxes2 (Tensor[M, 4])
    Returns:
        iou (Tensor[N, M]): the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2
    """
    area1 = box_area(boxes1)  # 每个框的面积 (N,)
    area2 = box_area(boxes2)  # (M,)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2] # N中一个和M个比较; 所以由N,M 个
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]  #小于0的为0  clamp 钳;夹钳;
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]

    iou = inter / (area1[:, None] + area2 - inter)
    return iou  # NxM, boxes1中每个框和boxes2中每个框的IoU值;


def nms(boxes: Tensor, scores: Tensor, iou_threshold: float):
    """
    :param boxes: [N, 4], 此处传进来的框,是经过筛选(NMS之前选取过得分TopK)之后, 在传入之前处理好的;
    :param scores: [N]
    :param iou_threshold: 0.7
    :return:
    """
    keep = []  # 最终保留的结果, 在boxes中对应的索引;
    idxs = scores.argsort()  # 值从小到大的 索引
    while idxs.numel() > 0:  # 循环直到null; numel(): 数组元素个数
        # 得分最大框对应的索引, 以及对应的坐标
        max_score_index = idxs[-1]
        max_score_box = boxes[max_score_index][None, :]  # [1, 4]
        keep.append(max_score_index)
        if idxs.size(0) == 1:  # 就剩余一个框了;
            break
        idxs = idxs[:-1]  # 将得分最大框 从索引中删除; 剩余索引对应的框 和 得分最大框 计算IoU;
        other_boxes = boxes[idxs]  # [?, 4]
        ious = box_iou(max_score_box, other_boxes)  # 一个框和其余框比较 1XM
        idxs = idxs[ious[0] <= iou_threshold]

    keep = idxs.new(keep)  # Tensor
    return keep

>>> 如有疑问,欢迎评论区一起讨论

GitHub 加速计划 / vi / vision
15.85 K
6.89 K
下载
pytorch/vision: 一个基于 PyTorch 的计算机视觉库,提供了各种计算机视觉算法和工具,适合用于实现计算机视觉应用程序。
最近提交(Master分支:3 个月前 )
518ee93d 7 天前
7d077f13 9 天前
Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐