YOLOv5推理详解及预处理高性能实现

注意事项

一、2022/12/6更新

将原文中Latex编辑的公式替换成图片。CSDN公式渲染的效果真是一言难尽😂

二、2022/11/30更新

更新详解YOLOv5预处理和后处理过程,利用CUDA核函数实现高性能预处理过程

看有小伙伴收藏了,看来得把欠下的补上了😄。加更!加更!加更!

最近刚好有时间,将YOLOv5的预处理和后处理流程完整的梳理一遍吧。

前言

梳理下YOLOv5预处理和后处理流程,并实现CUDA核函数版预处理(以jetson nano作为本次实验设备)

视频讲解:bilibili/手写AI/预处理高性能抖音/手写AI/仿射变换cuda例子

代码参考:https://github.com/shouxieai/tensorRT_Pro/blob/main/src/tensorRT/common/preprocess_kernel.cu

1.YOLOv5预处理

1.1 预处理分析

这里只做简单分析,更多细节请观看预处理高性能视频

首先来看yolov5预处理部分,在目标检测中,我们的预处理通常是先对图像进行等比缩放,然后居中,多余部分填充,类似于下图所展示的,这个过程也被叫做添加灰条(即letterbox)

在这里插入图片描述

图(1) 预处理

整个过程可以分为三个步骤

1.等比缩放,矩阵YOLOv5推理详解及预处理高性能实现实现(缩放倍数Scale为目标图像与源图像宽比值和高比值的最小值)

在这里插入图片描述

图(2) 缩放S

2.将图片中心平移到左上角坐标原点,矩阵YOLOv5推理详解及预处理高性能实现实现

在这里插入图片描述

图(3) 平移O

3.将图片平移到目标位置的中心,矩阵YOLOv5推理详解及预处理高性能实现实现

在这里插入图片描述

图(4) 平移T

预处理的过程可以通过三个矩阵完成,分别为缩放矩阵YOLOv5推理详解及预处理高性能实现,平移矩阵YOLOv5推理详解及预处理高性能实现,平移矩阵YOLOv5推理详解及预处理高性能实现,将这三个矩阵可以进行合并成一个矩阵YOLOv5推理详解及预处理高性能实现,其中YOLOv5推理详解及预处理高性能实现矩阵被称为仿射变换矩阵,可以帮助我们完成图像预处理工作,后续的CUDA核函数就是基于YOLOv5推理详解及预处理高性能实现矩阵去进行加速的

我们可以直接写出仿射变换矩阵YOLOv5推理详解及预处理高性能实现,如下式(4)所示,其中代YOLOv5推理详解及预处理高性能实现表源图像,YOLOv5推理详解及预处理高性能实现代表目标图像,YOLOv5推理详解及预处理高性能实现为源图像上任意像素点坐标,YOLOv5推理详解及预处理高性能实现为目标图像上任意像素点坐标。

在这里插入图片描述

通过上述变换可以得到源图像任意像素值的坐标对应在目标图像的位置,那么整个预处理过程也就可以通过仿射变换矩阵YOLOv5推理详解及预处理高性能实现来完成。除了将源图像通过YOLOv5推理详解及预处理高性能实现变换到目标图像,我们也要考虑将目标图像映射回源图像,即求YOLOv5推理详解及预处理高性能实现矩阵的逆矩阵YOLOv5推理详解及预处理高性能实现

为什么要求逆变换呢?主要是因为网络推理时是在目标图像上进行的,换而言之预测框是绘制在目标图像上的,而最终我们所需要的是源图像的预测框,从目标图像的预测框到源图像预测框的转换就是通过逆矩阵YOLOv5推理详解及预处理高性能实现完成的

逆变换矩阵YOLOv5推理详解及预处理高性能实现如下


在这里插入图片描述

则有在这里插入图片描述

至此,预处理分析讲解完成

1.2 仿射变换正逆变换实验

既然知道预处理可以通过仿射变换矩阵YOLOv5推理详解及预处理高性能实现来完成,现在来进行一个小实验。读取一张图片,通过YOLOv5推理详解及预处理高性能实现矩阵完成预处理,然后通过逆变换矩阵YOLOv5推理详解及预处理高性能实现恢复原图。

代码如下

import numpy as np
import cv2
import matplotlib.pyplot as plt

def inv_align(M):
    # M的逆变换
    k  = M[0, 0]
    b1 = M[0, 2]
    b2 = M[1, 2]
    return np.array([
        [1/k, 0, -b1/k],
        [0, 1/k, -b2/k]
    ])

def align(image, dst_size):
    # image->[h,w,c]
    oh, ow = image.shape[:2]
    dh, dw = dst_size
    scale = min(dw/ow, dh/oh)
    # 仿射变换矩阵M
    M = np.array([
        [scale, 0, -scale * ow * 0.5 + dw * 0.5],
        [0, scale, -scale * oh * 0.5 + dh * 0.5]
    ], dtype=np.float32)
    # warpAffine之后的图片
    image_pre = cv2.warpAffine(image, M, dst_size, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(114,114,114))
    return image_pre, M, inv_align(M)

if __name__ == '__main__':
    cat = cv2.imread("cat.png")
    cat_pre, M, IM = align(cat, (640,640))
    # 将图片进行warpaffine
    plt.subplot(1, 2, 1)
    plt.title("WarpAffine")
    plt.imshow(cat_pre[..., ::-1])
    # 将warpaffine后的图片复原
    cat_res = cv2.warpAffine(cat_pre, IM, cat.shape[:2][::-1])
    plt.subplot(1, 2, 2)
    plt.title("Restore")
    plt.imshow(cat_res[...,::-1])
    plt.show()

在这里插入图片描述

图(5) 仿射变换实验效果

1.3 YOLOv5的预处理过程

YOLOv5的预处理部分除了上述仿射变换到一定尺寸大小的图片之外,输入到网络中的实际上是一个4个维度的Tensor(B,C,H,W),所以还需要进行如下变换:

  • BGR ➡ RGB(以opencv读取图片为例)
  • /255.0
  • 通道数变换 H,W,C ➡ C,H,W(以opencv读取图片为例)
  • 添加batch维度 C,H,W ➡ B,C,H,W

预处理代码如下

import numpy as np
import cv2

def preprocess(img, dst_width=640, dst_height=640):
    '''
    :param img: 输入的图片
    :param dst_width: 预处理后的图像宽
    :param dst_height: 预处理后的图像高
    :return: 预处理后的图片,仿射变换矩阵的逆变换矩阵IM
    '''
    scale = min((dst_width / img.shape[1]), (dst_height / img.shape[0]))
    ox = (-scale * img.shape[1] + dst_width)  / 2
    oy = (-scale * img.shape[0] + dst_height) / 2
    M  = np.array([
        [scale, 0, ox],
        [0, scale, oy]
    ], dtype=np.float32)
    # img_pre为仿射变换后的图即原始图像缩放到[dst_width,dst_height]
    img_pre = cv2.warpAffine(img, M, dsize=[dst_width, dst_height], flags=cv2.INTER_LINEAR,
                             borderMode=cv2.BORDER_CONSTANT, borderValue=(114,114,114))
    # cv2.imshow("img_pre", img_pre)
    # cv2.waitKey(0)
    IM = cv2.invertAffineTransform(M)
    # -----------------------------------------------------------------------#
    #   需要进行的预处理
    #   1. BGR -> RGB
    #   2. /255.0
    #   3. 通道数变换 H,W,C -> C,H,W
    #   4. 添加batch维度 C,H,W -> B,C,H,W
    # -----------------------------------------------------------------------#
    img_pre = (img_pre[...,::-1] / 255.0).astype(np.float32)
    img_pre = img_pre.transpose(2,0,1)[None]
    return img_pre, IM

YOLOv5整个预处理可以通过一个preprocess()函数完成,输入是opencv读取的图片以及仿射变换目标图片的宽高,输出是预处理后的img_pre(4维Tensor)以及仿射变换的逆矩阵YOLOv5推理详解及预处理高性能实现。我们只需要将预处理后的img_pre(4维Tensor)放入网络中进行推理得到预测结果,然后利用逆矩阵YOLOv5推理详解及预处理高性能实现获得我们最终的结果。

至此,YOLOv5预处理讲解完成。

2 双线性插值

2.1 What

什么是双线性插值?

图像缩放是图像处理中的一个特别基础的操作,其中用到的就是插值算法

简单来说,插值指利用已知的点来"猜"未知的点,图像领域插值常用在修改图像尺寸的过程中,由旧的图像矩阵中的点计算新图像矩阵中的点并插入,不同的计算过程就是不同的插值算法。

常见的插值算法有:最近邻法(Nearest Interpolation)、双线性插值(Bilinear Interpolation)、双三次插值(Bicubic Interpolation)

参考自一篇文章为你讲透双线性插值

2.2 Why

为什么要了解双线性插值?

在上述YOLOv5预处理过程中,我们求出仿射变换矩阵YOLOv5推理详解及预处理高性能实现后调用cv2.warpAffine()函数便可实现图像缩放、居中、多余部分填充。目前预处理是调用第三方库函数实现的,那么当我们需要编写CUDA核函数来加速运行时,需要自己手动实现cv2.warpAffine()这个过程,那么我们需要了解cv2.warpAffine具体是怎么实现的,实现的原理是什么?cv2.warpAffine()实现的算法就是刚才提到的插值算法

2.3 How

如何去实现双线性插值?

首先来看下cv2.warpAffine()函数的参数

  • img—待处理的图片
  • M—仿射变换矩阵
  • dsize—目标图像的尺寸,YOLOv5为[640,640]
  • flags—插值算法,如cv2.INTER_LINEAR(双线性插值)
  • borderMode—边界填充模式,如cv2.BORDED_CONSTANT(常量填充)
  • boderValue—常量填充模式下填充的常量值,如borderValue=(114,114,114)

知道具体参数后,我们该怎么实现呢?在前面我们通过仿射变换矩阵M可以从源图像的任意一点像素得到目标图像的任意像素,现在我们未知的是目标图像,即目标图像的每个像素值该如何取舍,这样就要利用到逆矩阵YOLOv5推理详解及预处理高性能实现,去取源图像的像素然后填充目标图像,如下图所示

在这里插入图片描述

图(6) 插值过程

上述图(a)➡图(c)便是整个预处理过程,中间通过图(b)实现,首先准备[640,640]大小的空白图像(以YOLOv5为例)即图(b),现在要往图(b)填充像素使其最终得到图(c),这是我们的目的。对于图(b)中像素点为(100,200)的位置该填充什么像素呢?我们可以通过逆矩阵YOLOv5推理详解及预处理高性能实现计算得到对应的(100,200)的像素位置对应于原图的(60.8,70.2)的像素位置,但是像素都是整数,那么小数位置(60.8,70.2)的像素该如何取呢?这时候就需要使用到插值算法,这里使用双线性插值,具体实现见下图

在这里插入图片描述

图(7) 双线性插值

假设要获取未知像素点(60.8,70.2)的像素值,即图中红色区域的像素值,我们可以先找其周围最近的四个像素点P1(60,70)、P2(61,70)、P3(61,71)、P4(60,71),这四个像素点的像素值是已知的,然后再计算这四个像素点到未知像素点对应的面积作为其权重即W1、W2、W3、W4,那么最终目标图像(100,200)位置的像素值可用下式计算

在这里插入图片描述

目标图像[100,200]位置的像素值便可通过上式计算得出,同理对于其它位置的像素点可用相同的方法,那么循环整幅图像后即可得到每个像素点的像素值,填充完成后便可得到图(c),特别注意的是,通过逆矩阵计算源图像位置时可能会超出源图像位置,这时会找不到对应的像素值,此时在对应的目标图像上填充常量像素值即可,也就是图(c)中的灰条位置。

至此,双线性插值部分讲解完成。接下来看个双线性插值的小实验。

2.4 双线性插值实验

这里手动来实现一个warpAffine,并对比cv2.warpAffine看各自的效果

代码如下

import cv2
import numpy as np
import matplotlib.pyplot as plt

def inv_align(M):
    # M的逆变换
    k = M[0, 0]
    b1 = M[0, 2]
    b2 = M[1, 2]
    return np.array([
        [1/k, 0, -b1/k],
        [0, 1/k, -b2/k]
    ])

def align(image, dst_size):
    # image->[h,w,c]
    oh, ow = image.shape[:2]
    dh, dw = dst_size
    scale = min(dw/ow, dh/oh)
    # 仿射变换矩阵M 
    M = np.array([
        [scale, 0, -scale * ow / 2 + dw / 2],
        [0, scale, -scale * oh / 2 + dh / 2]
    ], dtype=np.float32)
    return M

def pyWarpAffine(image, M, dst_size, constant=(0, 0, 0)):
    
    # 注意输入的M矩阵格式,是Origin->Dst
    # 而这里需要的是Dst->Origin,所以要取逆矩阵
    # M = cv2.invertAffineTransform(M)
    IM = inv_align(M)
    constat = np.array(constant)
    ih, iw  = image.shape[:2]       # 源图像宽高
    dw, dh  = dst_size              # 目标图像宽高
    dst     = np.full((dh, dw, 3), constant, dtype=np.uint8)    # 空白图像,全填充常量
    irange  = lambda p : p[0] >= 0 and p[0] < iw and p[1] >= 0 and p[1] < ih    # 匿名函数用于边界判断
    
    for y in range(dh):
        for x in range(dw):

            homogeneous = np.array([[x, y, 1]]).T       # 目标图像任意像素点
            ox, oy = IM @ homogeneous                   # 对应源图像像素点
            
            low_ox = int(np.floor(ox))
            low_oy = int(np.floor(oy))
            high_ox = low_ox + 1
            high_oy = low_oy + 1

            # p1    p2
            #    o
            # p4    p3

            pos = ox - low_ox, oy - low_oy
            p1_area = (1 - pos[0]) * (1 - pos[1])
            p2_area = pos[0] * (1 - pos[1])
            p3_area = pos[0] * pos[1]
            p4_area = (1 - pos[0]) * pos[1]

            p1 = low_ox, low_oy
            p2 = high_ox, low_oy
            p3 = high_ox, high_oy
            p4 = low_ox, high_oy
            p1_value = image[p1[1], p1[0]] if irange(p1) else constant
            p2_value = image[p2[1], p2[0]] if irange(p2) else constant
            p3_value = image[p3[1], p3[0]] if irange(p3) else constant
            p4_value = image[p4[1], p4[0]] if irange(p4) else constant
            dst[y, x] = p1_area * p1_value + p2_area * p2_value + p3_area * p3_value + p4_area * p4_value
            # 后续可以完成更多预处理操作,达到像素级预处理
            # 交换bgr rgb
            # normalize -> mean / std
            # 1行代码实现normalize -> /255.0
            # bgr bgr bgr -> bbb ggg rrr
    return dst

if __name__ == "__main__":
    cat = cv2.imread("cat1.png")
    M = align(cat, [640, 640])
    
    cat_cv = cv2.warpAffine(cat, M, [640, 640])
    plt.subplot(1, 2, 1)
    plt.title("OpenCV")
    plt.imshow(cat_cv[..., ::-1])
    
    cat_py = pyWarpAffine(cat, M, [640, 640])
    plt.subplot(1, 2, 2)
    plt.title("PyWarpAffine")
    plt.imshow(cat_py[..., ::-1])
    # plt.savefig("./warpAffinn.png", bbox_inches='tight')
    plt.show()

在这里插入图片描述

图(8) 双线性插值对比

3.YOLOv5后处理

YOLOv5的后处理主要包括预测框的decode解码和NMS非极大值抑制两部分,对于一张640×640图片来说,YOLOv5预测框的总数量是25200(针对COCO数据集80类别而言),每个预测框的维度是85
YOLOv5推理详解及预处理高性能实现
其中的5分别对应的是cx,cy,w,h,obj_conf,分别代表的含义是边界框中心点坐标、宽高、边界框内包含物体的置信度。80对应的是COCO数据中的80个类别。

为什么需要decode解码呢?

很简单,你推理的图片是在预处理后的图片上的,框都是画在预处理完成后的图片上的(即640×640的目标图像),而你实际想要的源图像的一个效果图,怎么将目标图像上的框映射回源图像上呢?大家可能想到的就是通过我们之前的仿射变换逆矩阵YOLOv5推理详解及预处理高性能实现,同时为了方便后续的NMS非极大值抑制,我们需要将预测框的属性进行变换,从中心点宽高➡左上角右下角,那么整个decode解码过程就完成了。

代码如下

def postprocess(pred, IM, conf_thres=0.25):
    # 输入是模型推理的结果,即25200个预测框
    # 1,25200,85 [cx,cy,w,h,obj_conf,cls]
    boxes = []
    for img_id, box_id in zip(*np.where(pred[...,4] > conf_thres)):
        item = pred[img_id, box_id]
        cx, cy, w, h, obj_conf = item[:5]
        label = item[5:].argmax()
        confidence = obj_conf * item[5 + label]
        if confidence < conf_thres:
            continue
        left  = cx - w * 0.5
        top = cy - h* 0.5
        right = cx + w * 0.5
        bottom  = cy + h * 0.5
        boxes.append([left, top, right, bottom, confidence, img_id, label])

    boxes = np.array(boxes)
    lr = boxes[:,[0, 2]]
    tb = boxes[:,[1, 3]]
    boxes[:,[0,2]] = IM[0][0] * lr + IM[0][2]
    boxes[:,[1,3]] = IM[1][1] * tb + IM[1][2]
    boxes = sorted(boxes.tolist(), key=lambda x:x[4], reverse=True)
    
    return boxes

为什么需要NMS非极大值抑制呢?

因为模型推理出来的预测框很多都是重叠的,我们需要选择最优的,具体怎么选呢?那就是通过NMS算法,关于NMS算法这里就不再赘述了,大家可以自行学习,网上示例代码也很多。这里给出杜老师写过的一个NMS实现吧,供大家参考。

def boxes_iou(box1, box2):
    def box_area(box):
        return (box[2] - box[0]) * (box[3] - box[1])
    # left, top     = max(box1[:2], box2[:2])
    # right, bottom = min(box1[2:4], box2[2:4])
    left   = max(box1[0], box2[0])
    top    = max(box1[1], box2[1])
    right  = min(box1[2], box2[2])
    bottom = min(box1[3], box2[3])
    cross = max(0, right-left) * max(0, bottom-top)
    union = box_area(box1) + box_area(box2) - cross
    if cross == 0 or union == 0:
        return 0
    return cross / union

def nms(boxes, iou_thres=0.5):
    keep_boxes = []
    remove_flags = [False] * len(boxes)
    for i in range(len(boxes)):
        if remove_flags[i]:
            continue
        ibox = boxes[i]
        keep_boxes.append(ibox)
        for j in range(len(boxes)):
            if remove_flags[j]:
                continue
            jbox = boxes[j]
            if ibox[5] != jbox[5] or ibox[6] != jbox[6]:
                continue
            iou = boxes_iou(ibox, jbox)
            if iou >= iou_thres:
                remove_flags[j] = True
    return keep_boxes

至此,YOLOv5后处理部分完毕。

4.YOLOv5推理

通过对YOLOv5的预处理和后处理讨论完成之后,整个的推理过程就显而易见了。YOLOv5推理包括预处理和后处理两部分,其中预处理主要包括warpAffine和双线性插值,后处理主要包括decode解码和NMS两部分。

废话少说直接上代码,提供代码权重图片下载链接Biadu Drive【password:yolo】,大家可自行下载测试

import numpy as np
import cv2
import onnxruntime

def preprocess(img, dst_width=640, dst_height=640):
    '''
    :param img: 输入的图片
    :param dst_width: 预处理后的图像宽
    :param dst_height: 预处理后的图像高
    :return: 预处理后的图片,仿射变换矩阵的逆变换矩阵IM
    '''
    scale = min((dst_width / img.shape[1]), (dst_height / img.shape[0]))
    ox = (-scale * img.shape[1] + dst_width)  / 2
    oy = (-scale * img.shape[0] + dst_height) / 2
    M  = np.array([
        [scale, 0, ox],
        [0, scale, oy]
    ], dtype=np.float32)
    # img_pre为仿射变换后的图即原始图像缩放到[dst_width,dst_height]
    img_pre = cv2.warpAffine(img, M, dsize=[dst_width, dst_height], flags=cv2.INTER_LINEAR,
                             borderMode=cv2.BORDER_CONSTANT, borderValue=(114,114,114))
    IM = cv2.invertAffineTransform(M)
    # -----------------------------------------------------------------------#
    #   需要进行的预处理
    #   1. BGR -> RGB
    #   2. /255.0
    #   3. 通道数变换 H,W,C -> C,H,W
    #   4. 添加batch维度 C,H,W -> B,C,H,W
    # -----------------------------------------------------------------------#
    img_pre = (img_pre[...,::-1] / 255.0).astype(np.float32)
    img_pre = img_pre.transpose(2,0,1)[None]
    return img_pre, IM

def iou(box1, box2):
    def area_box(box):
        return (box[2] - box[0]) * (box[3] - box[1])
    # box -> [x1,y1,x2,y2,...]
    left, top = max(box1[:2], box2[:2])
    right, bottom = min(box1[2:4], box2[2:4])
    union = max((right-left), 0) * max((bottom-top), 0)
    cross = area_box(box1) + area_box(box2) - union
    if cross == 0 or union == 0:
        return 0
    return union / cross

def NMS(boxes, iou_thresh=0.45):
    '''
    :param boxes: decode解码排序后的boxes [n,7] 7 = x1,y1,x2,y2,conf,img_id,label
    :param iou_thresh: iou阈值
    :return: 经过NMS的boxes
    '''
    # 利用remove_flags标记需要去除的box
    remove_flags = [False] * len(boxes)
    # 保留下的box
    keep_boxes = []
    for i in range(len(boxes)):
        if remove_flags[i]:
            continue
        ibox = boxes[i]
        keep_boxes.append(ibox)
        for j in range(len(boxes)):
            if remove_flags[j]:
                continue
            jbox = boxes[j]
            # 只有同一张图片中的同一个类别的box才计算iou
            if ibox[5] != jbox[5] or ibox[6] != jbox[6]:
                continue
            # 计算iou,若大于阈值则标记去除
            if iou(ibox, jbox) > iou_thresh:
                remove_flags[j] = True
    return keep_boxes


def postprocess(pred, IM, iou_thresh=0.45, conf_thresh=0.25):
    '''
    :param pred: 模型推理的结果 [1,25200,85] 85 = cx,cy,w,h,conf + 80
    :param IM: 仿射变换矩阵的逆变换,主要用来将box映射回原图
    :param iou_thresh: iou阈值
    :param cof_thresh: 置信度阈值
    :return: 经过NMS的boxes
    '''
    # 保存decode解码后的boxes
    boxes = []
    for img_id, box_id in zip(*np.where(pred[...,4] >= conf_thresh)):
        item = pred[img_id][box_id]
        cx, cy, w, h, obj_conf = item[:5]
        label = item[5:].argmax()
        confidence = obj_conf * item[5+label]
        if confidence < conf_thresh:
            continue
        left = cx  - w * 0.5
        top  = cy  - h * 0.5
        right  = cx  + w * 0.5
        bottom = cy  + h * 0.5
        boxes.append([left, top, right, bottom, confidence, img_id, label])
    # 利用IM将box映射回原图
    boxes = np.array(boxes)
    lr = boxes[..., [0, 2]]
    tb = boxes[..., [1, 3]]
    boxes[..., [0, 2]] = lr * IM[0][0] + IM[0][2]
    boxes[..., [1, 3]] = tb * IM[1][1] + IM[1][2]
    # 将boxes按照置信度高低排序
    boxes = sorted(boxes.tolist(), key= lambda x : x[4], reverse=True)
    # 将排序后的boxes作NMS
    return NMS(boxes, iou_thresh=iou_thresh)

if __name__ == '__main__':
    img = cv2.imread("bus.jpg")
    # 预处理
    img_pre, IM = preprocess(img)
    session = onnxruntime.InferenceSession("yolov5s.onnx", providers=["CUDAExecutionProvider"])
    # 模型推理
    pred = session.run(["output"], {"images" : img_pre})[0]
    # 后处理
    boxes = postprocess(pred, IM)
    for obj in boxes:
        x1, y1, x2, y2 = map(int, obj[:4])
        label = int(obj[6])
        confidence = obj[4]
        cv2.rectangle(img, (x1, y1), (x2, y2), (0,255,0), 1, 8)
        cv2.putText(img, f"{label}:{confidence:.3f}", (x1, y1-6), 0, 1, (0,0,255), 2, 8)
    cv2.imshow("img_pre", img)
    cv2.waitKey(0)

效果如下图
在这里插入图片描述

图(9) YOLOv5推理效果图

在这里为了方便演示使用onnx来进行推理,目的是让大家了解YOLOv5整个推理流程,方便大家去编写CUDA核函数以实现加速的目的。

5.高性能预处理-CUDA核函数实现

之前和大家分享了YOLOv5预处理可通过warpAffine和双线性插值来实现,现在我们需要编写CUDA核函数来加速预处理过程,首先我们先简单聊聊CUDA相关知识。

5.1 CUDA相关

关于CUDA博主也不太了解,需要各位自行学习,推荐两个官方文档CUDA C++ Best Practices GuideCUDA C++ Programming Guide。这里也找了一篇不错的文章分享,内容copy,强烈建议读原文

copy自:CUDA编程入门极简教程

作者:小小将

5.1.1 前言

先来关注两个问题,CUDA是什么?CUDA与cuDNN关系?

CUDA是什么?

CUDA(Compute Unified Device Architecture)是建立在NVIDIA的GPU上的一个通用并行计算平台和编程模型,基于CUDA编程可以利用GPU的并行计算引擎解决复杂的计算

CUDA与cuDNN关系?

首先,CUDA是C语言在GPU编程上的拓展包,cuDNN(CUDA® Deep Neural Network library)是封装了深度学习算子(如卷积)的库

关系:CUDA可以用来实现cuDNN定义的各种接口,早期cuDNN应该内部使用CUDA实现的,但随着NVIDIA软件生态的发展,cuDNN选择更底层,更靠近硬件,更难用的工具来构建Kernel(核函数),比如PTX,比如直接写汇编(SASS)

NVDIA依靠CUDA确立了高性能计算尤其是神经网络高性能计算的地位,CUDA肩负着软件生态通用性的任务,而高性能的任务,更多需要cuDNN、cuBLAS(CUDA® Basic Linear Algebra Subroutine library)高性能软件来承担。成熟的算子如conv、fc,用户使用库可以直接获得最优性能,而对新的算子或者用户独有的算子,可以通过CUDA简单实现一个性能还可以接受的版本。最后通过TensorRT、Pytorch这样的框架将二者衔接起来。copy自如何评价cuDNN和CUDA的关系?

CUDA并不是一个独立运行的计算平台,而需要与CPU协同工作。GPU并行计算其实指的是基于CPU+GPU的异构计算架构。在异构计算架构中,GPU与CPU通过PCIe总线连接在一起来协同工作,CPU所在位置称为主机端(Host),GPU所在位置称为设备端(Device),如下图所示

在这里插入图片描述

图(10) 基于CPU+GPU的异构计算

5.1.2 CUDA编程模型基础

hostdeviceCUDA中的两个重要概念,host指代CPU及其内存,device指代GPU及其内存。CUDA程序中既包含host程序又包含device程序,它们分别在CPU和GPU上运行。同时,host与device之间可以进行通信,可以进行数据的拷贝。典型的CUDA程序的执行流程如下:

  • 1.分配host内存,并进行数据的初始化
  • 2.分配device内存,并从host将数据拷贝到device上
  • 3.调用CUDA的核函数在device上完成指定的计算
  • 4.将device上的运算结果拷贝到host上
  • 5.释放device和host上分配的内存

上面流程中最重要的一个过程是调用CUDA的核函数(kernel)来执行并行计算,kernel是CUDA编程中一个重要的概念,指在device上线程并行执行的函数,核函数用__global__符号声明,在调用时需要用<<<grid, block>>>来指定kernel要执行的线程数量,在CUDA中,每一个线程都要执行核函数,并且每个线程会分配一个唯一的线程号thread ID,这个ID值可以通过核函数内置变量threadIdx来获得。

由于GPU实际上是异构模型,所有需要区分host和device上的代码,在CUDA中是通过函数类型限定词区分host和device上的函数,主要由三个函数类型限定词如下:

  • __global__表示核函数,在device上之执行,从host中调用(一些特定的GPU也可以从device上调用),返回类型必须是void
  • __device__表示设备函数,在device上执行,仅可以从device上调用
  • __host__表示主机函数,仅可以从host上调用

要深刻理解kernel,必须要对kernel的线程层级结构有一个清晰的认识。首先GPU上有很多并行化的轻量级线程。kernel在device上执行时实际上是启动了很多线程,一个kernel所启动的所有线程称为一个网格(grid),同一个网格上的线程共享相同的全局内存空间,grid是线程结构的第一层次,而网格又可以分为很多线程块(block),一个线程块里面包含很多线程,这是第二个层次。线程两层组织结构如下图所示,这是一个grid和block均为2-dim的线程组织。grid和block都是定义在dim3类型的变量,dim3可以看成是包含三个无符号整数(x,y,z)成员的结构体变量,在定义时,缺省值初始化为1。因此grid和block可以灵活地定义为1-dim,2-dim以及3-dim结构,对于图中结构(水平方向为x轴),定义得grid和block如下所示,kernel在调用时也必须通过执行配置<<<grid,block>>>来指定kernel所使用得线程数及结构。

dim3 grid(3, 2);
dim3 block(5, 3);
kernel_func<<< grid, block >>>(params...);

在这里插入图片描述

图(11) Kernel上的两层线程组织结构(2-dim)

所以,一个线程需要两个内置的坐标变量(blockIdx,threadIdx)来唯一标识,它们都是dim3类型的变量,其中blockIdx指明线程所在grid中的位置,而threadIdx指明线程所在block中的位置,如图中的Thread(1,1)满足:

threadIdx.x = 1
threadIdx.y = 1
blockIdx.x = 1
blockIdx.y = 1

一个线程块(block)上的线程是放在同一个流式多处理器(streaming Multi-processor,SM)上的,但是单个SM的资源有限,这导致线程块(block)中的线程数是有限制的,现代GPUs的线程块可支持的线程数可达1024个。有时候,我们要知道一个线程在线程块(block)中的全局ID,此时就必须还要知道block的组织结构,这是通过线程的内置变量blockDim来获得。它获取线程块(block)各个维度的大小。对于一个2-dim的block(YOLOv5推理详解及预处理高性能实现),线程(YOLOv5推理详解及预处理高性能实现)的ID值为(YOLOv5推理详解及预处理高性能实现),如果是3-dim的block(YOLOv5推理详解及预处理高性能实现),线程(YOLOv5推理详解及预处理高性能实现)的ID值为(YOLOv5推理详解及预处理高性能实现)。另外线程还有内置变量gridDim,用于获得网格块各个维度的大小。

kernel的这种线程组织结构天然适合vector、matrix等运算,我们可以利用上图2-dim结构实现两个矩阵的加法,每个线程负责处理每个位置的两个元素相加,代码如下所示。线程块大小为(16,16),然后将N*N大小的矩阵均分为不同的线程块来执行加法运算。

// Kernel定义
__global__ void MatAdd(float A[N][N], float B[N][N], float C[N][N])
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;
    if (i < N && j < N)
        C[i][j] = A[i][j] + B[i][j];
}

int main()
{
    ...
    // Kernel 线程配置
    dim3 threadsPerBlock(16, 16);
    dim3 numBlocks(N / threadsPerBlock.x, N / threadsPerBlock.y);	// 向上取整
    // Kernel 调用
    MatAdd<<numBlocks, threadsPerBlock>>>(A, B, C);
    ...
}

此外这里简单聊下CUDA的内存模型,如下图所。可以看到,每个线程都有自己的私有本地内存(Local Memory),而每个线程块有共享内存(Shared Memory),可以被线程块(block)中所有线程共享,其生命周期与线程块一致。此外,所有的线程都可以访问全局内存(Global Memory)。还可以访问一些只读内存块:常量内存(Constant Memory)和纹理内存(Texture Memory)。更多的细节描述可观看视频CUDA编程之基础,内存模型和线程束

在这里插入图片描述

图(12) CUDA内存模型

kernel核函数实际上会启动很多线程,这些线程是逻辑上并行的,但是在物理层却不一定。GPU硬件的一个核心组件就是SM,当一个kernel核函数被执行时,它的grid中的线程块被分配到SM上,一个线程块只能在一个SM上被调度。而一个SM一般可以调度多个线程块,这要看SM本身能力。可以一个kernel的各个线程块被分配到多个SM,所以grid只是逻辑层,而SM才是执行的物理层。SM的基本执行单元是线程束(warps),线程束包含32个线程,但是一个SM同时并发的线程束数是有限的。总之,就是网格和线程块只是逻辑划分,一个kernel的所有线程其实在物理层不一定同时并发的。所以kernel的grid和block的配置不同,性能会出现差异。还有,由于SM的基本执行单元是包含32个线程的线程束,所以block大小一般要设置为32的倍数

在这里插入图片描述

图(13) CUDA编程的逻辑层和物理层

在进行CUDA编程前,可以先检查下jetson nano的硬件配置,这样才可以有的放矢,可以通过以下代码获得GPU的配置属性。

// main.cpp
#include <iostream>
#include <cuda_runtime_api.h>
#include <driver_types.h>

using namespace std;

int main()
{
    int dev = 0;
    cudaDeviceProp devProp;
    cudaGetDeviceProperties(&devProp, dev);
    std::cout << "使用GPU device" << dev << ": " << devProp.name << std::endl;
    std::cout << "SM的数量: " << devProp.multiProcessorCount << std::endl;
    std::cout << "每个线程块的共享内存大小: " << devProp.sharedMemPerBlock / 1024.0 << " KB" << std::endl;
    std::cout << "每个线程块的最大线程数: " << devProp.maxThreadsPerBlock << std::endl;
    std::cout << "每个EM的最大线程数: " << devProp.maxThreadsPerMultiProcessor << std::endl;
    std::cout << "每个SM的最大线程束数: " << devProp.maxThreadsPerMultiProcessor / 32 << std::endl;
    
    return 0; 
}

编译指令如下:

g++ main.cpp -o main -I /usr/local/cuda-10.2/include -L /usr/local/cuda-10.2/targets/aarch64-linux/lib/ -lcudart

输出如下图所示

在这里插入图片描述

图(14) Jetson nano的GPU配置属性

至此,CUDA相关的知识的介绍到这里就告一段落的,关于后续的向量加法实例和矩阵乘法实例大家可以自行查看原文CUDA编程入门极简教程

5.2 高性能预处理

相信大家通过CUDA相关知识的介绍应该对CUDA有了一定的了解。我们的目的是为了编写CUDA核函数实现warpAffine和双线性插值以达到高性能预处理。博主在这里对其做个简单介绍,更多细节可查看视频

视频讲解:抖音/手写AI/仿射变换CUDA例子

代码参考:https://github.com/shouxieai/tensorRT_Pro/blob/main/src/tensorRT/common/preprocess_kernel.cu

5.2.0 导言

神经网络模型的输入一般是固定大小的,比如640×640,但是我们都知道图片尺寸有大有小,怎么将图片大小转换成网络模型输入的大小呢?在这里可以通过warpaffine来实现,warpaffine的本质是对源图上的每一个像素点进行坐标变换,将源图上的像素点填充到目标图像上即可,而当目标图像通过坐标变换的逆变换即YOLOv5推理详解及预处理高性能实现去源图取值时,可能得出的坐标是非整数,无法进行取值,需要使用插值算法(如图(6)所示),整个预处理过程如下

在这里插入图片描述

图(15) 神经网络预处理过程

5.2.1 main函数

首先写一个main函数,用于实现总体框架,主要通过读取一张图像,经过预处理之后,保存下来,其代码如下

#include <cuda_runtime.h>
#include <opencv2/opencv.hpp>
#include <stdio.h>

using namespace cv;

#define min(a, b)   ((a) < (b) ? (a) : (b))
#define checkRuntime(op) __check_cuda_runtime((op), #op, __FILE__, __LINE__)

bool __check_cuda_runtime(cudaError_t code, const char* op, const char* file, int line){
    if(code != cudaSuccess){
        const char* err_name = cudaGetErrorName(code);
        const char* err_message = cudaGetErrorString(code);
        printf("runtime error %s:%d %s failed. \n code = %s, message = %s\n", file, line, op, err_name, err_message);
        return false;
    }
    return true;
}

int main(){
    Mat image = imread("cat.png");
    Mat output = warpaffine_to_center_align(image, Size(640, 640));
    imwrite("output.png", output);
    printf("Done");

    return 0;
}

5.2.2 warpaffine_to_center_align函数

该函数为host函数,用于分配内存,调用预处理函数等功能,首先来看该函数要实现的具体功能

  • 1.在CPU上开辟空间,用于保存预处理完成后的图像
  • 2.在GPU上开辟空间,将读取的图像数据拷贝到GPU上
  • 3.调用预处理函数
  • 4.将预处理后的图片重新拷贝到CPU上进行保存
  • 5.释放CPU和GPU上分配的内存

在这里插入图片描述

图(16) warpaffine_to_center_align函数功能

了解完warpaffine_to_center_align函数需要实现的功能之后,其代码就相对容易实现了,代码如下

void warp_affine_bilinear(// 声明
    uint8_t* src, int src_line_size, int src_width, int src_height,
    uint8_t* dst, int dst_line_size, int dst_width, int dst_height,
    uint8_t fill_value
);

Mat warpaffine_to_center_align(Mat image, const Size& size){
    // 步骤1 在CPU上开辟空间,用于保存预处理完成后的图像
    Mat output(size, CV_8UC3);

    // 步骤2 在GPU上开辟空间,将读取的图像数据拷贝到GPU上
    uint8_t* psrc_device = nullptr;
    uint8_t* pdst_device = nullptr;

    size_t src_size = image.cols * image.rows * 3;
    size_t dst_size = size.width * size.height * 3;

    checkRuntime(cudaMalloc(&psrc_device, src_size));
    checkRuntime(cudaMalloc(&pdst_device, dst_size));

    checkRuntime(cudaMemcpy(psrc_device, image.data, src_size, cudaMemcpyHostToDevice));
    
    // 步骤3 调用预处理函数
    warp_affine_bilinear(// 调用
        psrc_device, image.cols * 3, image.cols, image.rows,
        pdst_device, size.width * 3, size.width, size.height,
        114
    );

    // 步骤4 将预处理后的图片重新拷贝到CPU上进行保存
    checkRuntime(cudaPeekAtLastError());    // 检查核函数执行是否存在错误
    checkRuntime(cudaMemcpy(output.data, pdst_device, dst_size, cudaMemcpyDeviceToHost));

    // 步骤5 释放CPU和GPU上分配的内存
    checkRuntime(cudaFree(pdst_device));
    checkRuntime(cudaFree(psrc_device));

    return output;
}

5.2.3 warp_affine_bilinear函数

该函数为host函数,用于调用核函数,主要有以下几点说明

  • 1.首先确定grid_size和block_size的大小。我们启动的线程数为图像宽x图像高,每个线程处理一个像素(3通道),又由于之前提到过核函数启动时gird只是逻辑层,而SM才是执行的物理层,且SM的基本执行单元是包含32个线程的线程束,所以block一般设置为32的倍数,在这里block_size设置为2-dim的block(32, 32),grid_size可通过block_size的大小计算得到。
  • 2.定义仿射变化矩阵YOLOv5推理详解及预处理高性能实现
  • 3.调用warp_affine_bilinear_kernel核函数

代码如下

#include <cuda_runtime.h>
#include <iostream>

void warp_affine_bilinear(// 实现
    uint8_t* src, int src_line_size, int src_width, int src_height, // 像素0~255 uint8
    uint8_t* dst, int dst_line_size, int dst_width, int dst_height,
    uint8_t fill_value
){
    // 需要多少threads
    dim3 block_size(32, 32);
    dim3 grid_size((dst_width + 31) / 32, (dst_height + 31) / 32);
    
    AffineMatrix affine;
    affine.compute(Size(src_width, src_height), Size(dst_width, dst_height));

    warp_affine_bilinear_kernel<<<grid_size, block_size, 0, nullptr>>>(
        src, src_line_size, src_width, src_height,
        dst, dst_line_size, dst_width, dst_height,
        fill_value, affine
    );
}

5.2.4 warp_affine_bilinear_kernel核函数

核函数的功能主要是实现warpAffine和双线性插值即预处理,主要有以下几点说明

  • 1.一个thread负责一个像素(3通道),所以有全局Idx的计算
  • 2.已知源图像求目标图像的像素,通过仿射变换逆矩阵YOLOv5推理详解及预处理高性能实现可得,即代码中的affine_project()
  • 3.之前讨论过,目标图像通过YOLOv5推理详解及预处理高性能实现去取源图像时可能存在超出边界问题,此时填充常量值即可
  • 4.对于能取到值的像素,其映射回源图坐标可能是非整数,此时需要插值算法(双线性插值),如图(6)所示,计算距离最近的四个像素点的加权值作为最终的像素值,如图(7)所示
  • 5.已知全局Idx索引(dx,dy),如何确定像素填充的具体位置呢?可看图(16),图中是一个4×4大小的图片,如果我想要填充Idx=(1,2)的像素值,我如何获得其内存地址呢?假设图片起始地址为src,则有YOLOv5推理详解及预处理高性能实现

在这里插入图片描述

图(17) 由全局索引获取对应内存地址

代码如下

__device__ void affine_project(float* matrix, int x, int y, float* proj_x, float* proj_y){

    // matrix
    // m0, m1, m2
    // m3, m4, m5
    *proj_x = matrix[0] * x + matrix[1] * y + matrix[2];
    *proj_y = matrix[3] * x + matrix[4] * y + matrix[5];
}

__global__ void warp_affine_bilinear_kernel(
    uint8_t* src, int src_line_size, int src_width, int src_height,
    uint8_t* dst, int dst_line_size, int dst_width, int dst_height,
    uint8_t fill_value, AffineMatrix matrix
){
    // 全局idx的计算,一个thread负责一个像素(3通道)
    int dx = blockDim.x * blockIdx.x + threadIdx.x;
    int dy = blockDim.y * blockIdx.y + threadIdx.y;

    if(dx >= dst_width  || dy >= dst_height) return;
    
    float c0 = fill_value, c1 = fill_value, c2 = fill_value;
    float src_x = 0; float src_y = 0;
    affine_project(matrix.d2i, dx, dy, &src_x, &src_y);     // 由目标图像dx,dy以及M^-1计算源图像坐标src_x,src_y

    // 已知src_x,src_y
    if(src_x < -1 || src_x >= src_width || src_y < -1 || src_y >= src_height){
        // out of range
        // src_x < -1时,其高位high_x < 0,超出范围
        // src_x >= -1时,其高位high_x >= 0,存在取值
    }else{
        int y_low = floorf(src_y);
        int x_low = floorf(src_x);
        int y_high = y_low + 1;
        int x_high = x_low + 1;
        
        uint8_t const_values[] = {fill_value, fill_value, fill_value};
        float ly = src_y - y_low;
        float lx = src_x - x_low;
        float hy = 1 - ly;
        float hx = 1 - lx;
        float w1 = hy * hx,  w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
        uint8_t* v1 = const_values;
        uint8_t* v2 = const_values;
        uint8_t* v3 = const_values;
        uint8_t* v4 = const_values;
        if(y_low >= 0){
            if(x_low >= 0)
                v1 = src + y_low * src_line_size + x_low * 3;
            
            if(x_high < src_width)
                v2 = src + y_low * src_line_size + x_high * 3;
        }

        if(y_high < src_height){
            if(x_low >= 0)
                v3 = src + y_high * src_line_size + x_low * 3;
            
            if(x_high < src_width)
                v4 = src + y_high * src_line_size + x_high * 3;
        }

        c0 = floorf(w1 * v1[0] + w2 * v2[0] + w3 * v3[0] + w4 * v4[0] + 0.5f);
        c1 = floorf(w1 * v1[1] + w2 * v2[1] + w3 * v3[1] + w4 * v4[1] + 0.5f);
        c2 = floorf(w1 * v1[2] + w2 * v2[2] + w3 * v3[2] + w4 * v4[2] + 0.5f);
    }

    uint8_t* pdst = dst + dy * dst_line_size + dx * 3;
    pdst[0] = c0; pdst[1] = c1; pdst[2] = c2;
}

5.2.5 AffineMatrix结构体

结构体主要是为了计算仿射变换矩阵YOLOv5推理详解及预处理高性能实现以及逆矩阵YOLOv5推理详解及预处理高性能实现

代码如下

struct Size{
    int width = 0, height = 0;

    Size() = default;
    Size(int w, int h)
    :width(w), height(h){}
};

struct AffineMatrix{
    float i2d[6];       // image to dst(network), 2x3 matrix ==> M
    float d2i[6];       // dst to image, 2x3 matrix ==> M^-1

    // 求解M的逆矩阵,由于这个3x3矩阵的第三行是确定的0, 0, 1,因此可以简写如下
    void invertAffineTransform(float imat[6], float omat[6]){
        float i00 = imat[0]; float i01 = imat[1]; float i02 = imat[2];
        float i10 = imat[3]; float i11 = imat[4]; float i12 = imat[5];

        // 计算行列式
        float D = i00 * i11 - i01 * i10;
        D = D != 0 ? 1.0 / D : 0;
        
        // 计算剩余的伴随矩阵除以行列式
        float A11 = i11 * D;
        float A22 = i00 * D;
        float A12 = -i01 * D;
        float A21 = -i10 * D;
        float b1 = -A11 * i02 - A12 * i12;
        float b2 = -A21 * i02 - A22 * i12;
        omat[0] = A11; omat[1] = A12; omat[2] = b1;
        omat[3] = A21; omat[4] = A22; omat[5] = b2;
    }

    // 求解M矩阵
    void compute(const Size& from, const Size& to){
        float scale_x = to.width / (float)from.width;
        float scale_y = to.height / (float)from.height;

        float scale = min(scale_x, scale_y);
        /*
        M = [
        scale,    0,     -scale * from.width  * 0.5 + to.width  * 0.5
        0,     scale,    -scale * from.height * 0.5 + to.height * 0.5
        0,        0,                     1
        ]
        */
        /*
            - 0.5 的主要原因是使得中心更加对齐,下采样不明显,但是上采样时就比较明显
            参考:https://www.iteye.com/blog/handspeaker-1545126
        */
        i2d[0] = scale; i2d[1] = 0; i2d[2] = 
            -scale * from.width  * 0.5 + to.width  * 0.5 + scale * 0.5 - 0.5;
        
        i2d[3] = 0; i2d[4] = scale; i2d[5] =
            -scale * from.height * 0.5 + to.height * 0.5 + scale * 0.5 - 0.5;
        
        invertAffineTransform(i2d, d2i);
    }
};

5.3 完整示例代码及演示

大家可以直接下载博主提供好的示例代码BaiduDrive【password:yolo】,也可以参考下面完整的内容

main.cpp

#include <cuda_runtime.h>
#include <opencv2/opencv.hpp>
#include <stdio.h>
#include <algorithm>
#include <chrono>

using namespace cv;

#define min(a, b)   ((a) < (b) ? (a) : (b))
#define checkRuntime(op) __check_cuda_runtime((op), #op, __FILE__, __LINE__)

double timestamp_now_float() {
    return std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now().time_since_epoch()).count() / 1000.0;
}

bool __check_cuda_runtime(cudaError_t code, const char* op, const char* file, int line){
    if(code != cudaSuccess){
        const char* err_name = cudaGetErrorName(code);
        const char* err_message = cudaGetErrorString(code);
        printf("runtime error %s:%d %s failed. \n code = %s, message = %s\n", file, line, op, err_name, err_message);
        return false;
    }
    return true;
}

void warp_affine_bilinear(// 声明
    uint8_t* src, int src_line_size, int src_width, int src_height,
    uint8_t* dst, int dst_line_size, int dst_width, int dst_height,
    uint8_t fill_value
);

Mat warpaffine_to_center_align(Mat image, const Size& size){
    // 步骤1 在CPU上开辟空间,用于保存预处理完成后的图像
    Mat output(size, CV_8UC3);

    // 步骤2 在GPU上开辟空间,将读取的图像数据拷贝到GPU上
    uint8_t* psrc_device = nullptr;
    uint8_t* pdst_device = nullptr;

    size_t src_size = image.cols * image.rows * 3;
    size_t dst_size = size.width * size.height * 3;

    checkRuntime(cudaMalloc(&psrc_device, src_size));
    checkRuntime(cudaMalloc(&pdst_device, dst_size));

    checkRuntime(cudaMemcpy(psrc_device, image.data, src_size, cudaMemcpyHostToDevice));
    
    // 步骤3 调用预处理函数
    warp_affine_bilinear(// 调用
        psrc_device, image.cols * 3, image.cols, image.rows,
        pdst_device, size.width * 3, size.width, size.height,
        114
    );

    // 步骤4 将预处理后的图片重新拷贝到CPU上进行保存
    checkRuntime(cudaPeekAtLastError());    // 检查核函数执行是否存在错误
    checkRuntime(cudaMemcpy(output.data, pdst_device, dst_size, cudaMemcpyDeviceToHost));

    // 步骤5 释放CPU和GPU上分配的内存
    checkRuntime(cudaFree(pdst_device));
    checkRuntime(cudaFree(psrc_device));

    return output;
}

int main(){
    auto t0 = timestamp_now_float();
    Mat image = imread("cat.png");
    Mat output = warpaffine_to_center_align(image, Size(640, 640));
    imwrite("output.png", output);
    auto fee = timestamp_now_float() - t0;
    printf("Done. save to output.png fee = %.2f ms\n", fee);
    
    return 0;
}

affine.cu

#include <cuda_runtime.h>
#include <iostream>

#define min(a, b) ((a) < (b) ? (a) : (b))

struct Size{
    int width = 0, height = 0;

    Size() = default;
    Size(int w, int h)
    :width(w), height(h){}
};

struct AffineMatrix{
    float i2d[6];       // image to dst(network), 2x3 matrix ==> M
    float d2i[6];       // dst to image, 2x3 matrix ==> M^-1

    // 求解M的逆矩阵,由于这个3x3矩阵的第三行是确定的0, 0, 1,因此可以简写如下
    void invertAffineTransform(float imat[6], float omat[6]){
        float i00 = imat[0]; float i01 = imat[1]; float i02 = imat[2];
        float i10 = imat[3]; float i11 = imat[4]; float i12 = imat[5];

        // 计算行列式
        float D = i00 * i11 - i01 * i10;
        D = D != 0 ? 1.0 / D : 0;
        
        // 计算剩余的伴随矩阵除以行列式
        float A11 = i11 * D;
        float A22 = i00 * D;
        float A12 = -i01 * D;
        float A21 = -i10 * D;
        float b1 = -A11 * i02 - A12 * i12;
        float b2 = -A21 * i02 - A22 * i12;
        omat[0] = A11; omat[1] = A12; omat[2] = b1;
        omat[3] = A21; omat[4] = A22; omat[5] = b2;
    }

    // 求解M矩阵
    void compute(const Size& from, const Size& to){
        float scale_x = to.width / (float)from.width;
        float scale_y = to.height / (float)from.height;

        float scale = min(scale_x, scale_y);
        /*
        M = [
        scale,    0,     -scale * from.width  * 0.5 + to.width  * 0.5
        0,     scale,    -scale * from.height * 0.5 + to.height * 0.5
        0,        0,                     1
        ]
        */
        /*
            - 0.5 的主要原因是使得中心更加对齐,下采样不明显,但是上采样时就比较明显
            参考:https://www.iteye.com/blog/handspeaker-1545126
        */
        i2d[0] = scale; i2d[1] = 0; i2d[2] = 
            -scale * from.width  * 0.5 + to.width  * 0.5 + scale * 0.5 - 0.5;
        
        i2d[3] = 0; i2d[4] = scale; i2d[5] =
            -scale * from.height * 0.5 + to.height * 0.5 + scale * 0.5 - 0.5;
        
        invertAffineTransform(i2d, d2i);
    }
};

__device__ void affine_project(float* matrix, int x, int y, float* proj_x, float* proj_y){

    // matrix
    // m0, m1, m2
    // m3, m4, m5
    *proj_x = matrix[0] * x + matrix[1] * y + matrix[2];
    *proj_y = matrix[3] * x + matrix[4] * y + matrix[5];
}

__global__ void warp_affine_bilinear_kernel(
    uint8_t* src, int src_line_size, int src_width, int src_height,
    uint8_t* dst, int dst_line_size, int dst_width, int dst_height,
    uint8_t fill_value, AffineMatrix matrix
){
    // 全局idx的计算,一个thread负责一个像素(3通道)
    int dx = blockDim.x * blockIdx.x + threadIdx.x;
    int dy = blockDim.y * blockIdx.y + threadIdx.y;

    if(dx >= dst_width  || dy >= dst_height) return;
    
    float c0 = fill_value, c1 = fill_value, c2 = fill_value;
    float src_x = 0; float src_y = 0;
    affine_project(matrix.d2i, dx, dy, &src_x, &src_y);     // 由目标图像dx,dy以及M^-1计算源图像坐标src_x,src_y

    // 已知src_x,src_y
    if(src_x < -1 || src_x >= src_width || src_y < -1 || src_y >= src_height){
        // out of range
        // src_x < -1时,其高位high_x < 0,超出范围
        // src_x >= -1时,其高位high_x >= 0,存在取值
    }else{
        int y_low = floorf(src_y);
        int x_low = floorf(src_x);
        int y_high = y_low + 1;
        int x_high = x_low + 1;
        
        uint8_t const_values[] = {fill_value, fill_value, fill_value};
        float ly = src_y - y_low;
        float lx = src_x - x_low;
        float hy = 1 - ly;
        float hx = 1 - lx;
        float w1 = hy * hx,  w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
        uint8_t* v1 = const_values;
        uint8_t* v2 = const_values;
        uint8_t* v3 = const_values;
        uint8_t* v4 = const_values;
        if(y_low >= 0){
            if(x_low >= 0)
                v1 = src + y_low * src_line_size + x_low * 3;
            
            if(x_high < src_width)
                v2 = src + y_low * src_line_size + x_high * 3;
        }

        if(y_high < src_height){
            if(x_low >= 0)
                v3 = src + y_high * src_line_size + x_low * 3;
            
            if(x_high < src_width)
                v4 = src + y_high * src_line_size + x_high * 3;
        }

        c0 = floorf(w1 * v1[0] + w2 * v2[0] + w3 * v3[0] + w4 * v4[0] + 0.5f);
        c1 = floorf(w1 * v1[1] + w2 * v2[1] + w3 * v3[1] + w4 * v4[1] + 0.5f);
        c2 = floorf(w1 * v1[2] + w2 * v2[2] + w3 * v3[2] + w4 * v4[2] + 0.5f);
    }

    uint8_t* pdst = dst + dy * dst_line_size + dx * 3;
    pdst[0] = c0; pdst[1] = c1; pdst[2] = c2;
}

void warp_affine_bilinear(// 实现
    uint8_t* src, int src_line_size, int src_width, int src_height, // 像素0~255 uint8
    uint8_t* dst, int dst_line_size, int dst_width, int dst_height,
    uint8_t fill_value
){
    // 需要多少threads
    dim3 block_size(32, 32);
    dim3 grid_size((dst_width + 31) / 32, (dst_height + 31) / 32);
    
    AffineMatrix affine;
    affine.compute(Size(src_width, src_height), Size(dst_width, dst_height));

    warp_affine_bilinear_kernel<<<grid_size, block_size, 0, nullptr>>>(
        src, src_line_size, src_width, src_height,
        dst, dst_line_size, dst_width, dst_height,
        fill_value, affine
    );
}

CMakeLists.txt

cmake_minimum_required(VERSION 2.6)
project(pro)

option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
set(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/workspace)

# 如果你的opencv找不到,可以自己指定目录
# set(OpenCV_DIR   "/datav/lean/opencv-4.2.0/lib/cmake/opencv4/")

set(CUDA_TOOLKIT_ROOT_DIR     "/usr/local/cuda-10.2")

find_package(CUDA REQUIRED)
find_package(OpenCV)

include_directories(
    ${PROJECT_SOURCE_DIR}/src
    ${OpenCV_INCLUDE_DIRS}
    ${CUDA_TOOLKIT_ROOT_DIR}/include
    ${CUDNN_DIR}/include
)

link_directories(
    ${CUDA_TOOLKIT_ROOT_DIR}/lib64
    ${CUDNN_DIR}/lib
)

set(CMAKE_CXX_FLAGS  "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -O0 -Wfatal-errors -pthread -w -g")
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -std=c++11 -O0 -Xcompiler -fPIC -g -w ${CUDA_GEN_CODE}")
file(GLOB_RECURSE cpp_srcs ${PROJECT_SOURCE_DIR}/src/*.cpp)
file(GLOB_RECURSE cuda_srcs ${PROJECT_SOURCE_DIR}/src/*.cu)
cuda_add_library(plugin_list SHARED ${cuda_srcs})
target_link_libraries(plugin_list cuda cublas cudart cudnn)
target_link_libraries(plugin_list ${OpenCV_LIBS})

add_executable(pro ${cpp_srcs})

# 如果提示插件找不到,请使用dlopen(xxx.so, NOW)的方式手动加载可以解决插件找不到问题
target_link_libraries(pro cuda cublas cudart cudnn)
target_link_libraries(pro plugin_list)
target_link_libraries(pro ${OpenCV_LIBS})

完整的编译过程演示如下

在这里插入图片描述

图(18) 完整编译过程演示

运行完成后将会在workspace文件夹下生成预处理完成后的图片,如下图所示

在这里插入图片描述

图(19) 高性能预处理图

至此,预处理的CUDA核函数的讲解完成。在这里博主仅仅和大家做简单分析,具体更多的细节描述请看视频链接

结语

博主在这里针对YOLOv5的预处理和后处理做了简单的分析,同时与大家分享了预处理的高性能实现即CUDA核函数。目的是让大家理清思路,更好的完成后续的部署工作。

感谢各位看到最后,创作不易,读后有收获的看官请帮忙点个👍!!!

下载链接

  • YOLOv5推理代码权重图片下载链接Biadu Drive【password:yolo】
  • warpaffine示例代码下载链接BaiduDrive【password:yolo】
    • 只包含src、workspace、CMakeLists.txt三个文件
    • src下存放着源码main.cpp以及affine.cu
    • workspace存放需要预处理的图片以及可执行文件
    • CMakeList.txt以完成修改,可直接编译

参考

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
扎眼的阳光的头像扎眼的阳光普通用户
上一篇 2023年3月5日
下一篇 2023年3月5日

相关推荐

此站出售,如需请站内私信或者邮箱!