YOLOv5算法移植

最新推荐文章于 2024-04-30 07:43:30 发布

pssbjn

最新推荐文章于 2024-04-30 07:43:30 发布

阅读量1.4k

点赞数 38

文章标签： YOLO 算法

本文链接：https://blog.csdn.net/pssbjn/article/details/137782557

版权

YOLOv5算法移植

模型导出
裁剪
测试
模型转换
C/C++调用代码
- 转换代码

移植之前先要先创建数据集并训练得到权重文件(runs/train/expxx/weights/best.pt)。这里用缺省COCO 80分类的预训练模型作为示例。

模型导出

训练得到的PyTorch模型文件(.pt)仅包含权重文件，网络结构以python源代码索引的方式存在，不能直接用于移植，需要先转换为某种交换格式。多数带NPU(ASIC方案的加速器)的处理器的模型转换工具都支持caffe和onnx两种格式的模型文件。因为Caffe框架早已停止维护，如果网络结构中用到caffe不支持的算子，需要自行编写代码并编译实现特定的算子。转换caffe模型的好处是转换后的模型可以在PC上用C/C++代码加载和验证，方便估计量化导致的精度损失。

这里仅考虑支持onnx的情况。下载YOLOv5源代码，使用YOLOv5自带的export.py脚本导出ONNX模型文件。

python3 export.py --weights yolov5n.pt --include onnx --imgsz 416 736 --opset 12

参数说明:

weights: 训练得到的权重文件，这里用YOLOv5n预训练权重文件。
include: 导出格式
imgsz: 输入尺寸。
- 缺省的YOLOv5模型用三尺度检测，图像缩小为输入图像的1/32，因此输入尺寸必须是32的整数倍
- YOLOv5模型本身与输入图像的尺寸无关，图像输入模型之前会进行resize和padding，使得图像尺寸不超过最大尺寸(imgsz中的最大值)并保持目标的长宽比，同时处理后的图像尺寸是32的整数倍
- 嵌入式处理器中的resize和padding一般由硬件完成，并不一定保证能完成上述操作，一般仅resize，而不进行padding(padding除自身的内存搬移操作之外，还需要调用方记录偏移位置)。
- 多数CMOS摄像头的长宽比为16:9，因此选择将图像resize为736x416，同时满足长宽约比为16:9，长宽都为32的整数倍
opset: 算子版本号。根据模型转换工具支持的算子选择算子版本号。

导出完成后得到yolov5n.onnx，用netron可以查看网络结构。

裁剪

裁剪的目的是删除检测头Detect中不被模型转换工具支持的层或者不适合用NPU加速的层。

模型裁剪有三种方法:

用代码加载和修改onnx模型文件
用onnx-modifier可视化工具修改
修改导出代码

三种方法各有优缺点

用代码加载和修改onnx模型文件

先用netron查看网络结构，记住要删除或修改的层名称，对应的进行修改。代码如下:

import onnx
import onnxruntime
import re

input_path = 'yolov5n.onnx'
output_path = 'yolov5n-cut.onnx'

model = onnx.load(input_path)
graph = model.graph

# 删除detect层除conv和sigmoid之外的所有节点
conv, sigmod = []
for n in graph.node:
	if re.match(r'/model.24/m.\d/Conv', n.name):
		conv.append(n)
	elif elif n.name.startswith('/model.24/Sigmoid'):
		sigmoid.append(n)
	else:
		graph.node.remove(n)

# 连接conv层和sigmoid层
for c, s in zip(conv, sigmoid):
	s.input[0] = c.output[0]

# 删除所有输出，添加没有维度信息的临时输出
while len(graph.output) > 0: graph.output.pop()
for s in sigmoid: graph.output.extend([onnx.ValueInfoProto(name=s.output[0])])

# 计算输出shape
sess = onnxruntime.InferenceSession(model.SerializeToString())
inputs = sess.get_inputs()[0]
x = np.random.randn(*inputs.shape).astype(np.float32) # 随机输入
outputs = sess.run(None, {inputs.name: x})

# 删除所有输出
while len(graph.output) > 0: graph.output.pop()
# 重新添加输出
for s, o in zip(sigmoid, outputs):
    output = onnx.helper.make_tensor_value_info(s.output[0], onnx.TensorProto.FLOAT, o.shape)
    graph.output.extend([output, ])

# 保存裁剪后的模型文件
onnx.save(model, output_path)

这种方法一般要对每个模型文件编写一个裁剪文件，不太方便，在一个模型的训练优化验证过程中比较有用。

用onnx-modifier可视化工具修改、

是方法1的可视化版本，比较直观，操作有些繁琐。

修改导出代码

修改models/yolo.py中Detect类的forward函数

    def forward(self, x):
        z = []  # inference output
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv
            if self.export: z.append(x[i].sigmoid()) # 添加行
            ##### 以下不变 #######
            #####################
        if self.export: return z # 添加行

        return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)

修改代码后，执行export.py脚本得到的onnx文件是已经裁剪好的。这个方法对检测任务通用性比较好，不受Backbone结构的影响，修改两尺度或者四尺度输出也适用。缺点是分割任务的导出依赖检测任务的输出，修改检测任务的输出会导致分割任务的导出出错。这个问题可以通过增加export.py脚本参数的方式进行修正。

测试

裁剪完成之后验证裁剪后的模型是否可以正常加载和运行，基于PyTorch、Opencv和onnxruntime实现。代码如下：

加载模型文件、测试图片并运行

这一步与加载pt模型文件的流程基本一致，差别只在于运行结果preds不同

import onnxruntime as rt
import cv2
import torch.nn.functional as F

weights = 'yolov5n-cut.onnx'
image_file = 'test.jpg'

# 加载模型文件
sess = rt.InferenceSession(weights)
# 获取模型输入名称
input_name = sess.get_inputs()[0].name
# 获取模型输入尺寸
img_size = sess.get_inputs()[0].shape[-2:]、
# 从模型的元数据中获取类型名称(用于在图片上叠加类型信息)
classes = eval(sess.get_modelmeta().custom_metadata_map['names'])

# 加载图片并进行resize和padding
img = LoadImage(image_file, img_size)
# 运行模型
preds = sess.run(None, {input_name: img.astype(np.float32)})

后处理

后处理的目的是将运行模型得到的输出特征图转换为bouding box并进行NMS，得到最终结构。

2.1 获取预选框

预选框是从数据集中聚类得到的，每个尺度对应3个预选框。

yolov5_anchors = {
    8:  [10,13, 16,30, 33,23],       # P3/8
    16: [30,61, 62,45, 59,119],      # P4/16
    32: [116,90, 156,198, 373,326],  # P5/32
}

def yolov5cut_get_anchor(o, xs):
    c, h, w = o.shape
    assert c % 3 == 0, f'Invalid output channel number {c}, should be multiple of 3'
    assert c // 3 > (4 + 1), f'Invalid output channel number {c}, should > 15'
    assert xs[2] % h == 0, f'image height {xs[2]} is not multiple of output height {h}'
    assert xs[3] % w == 0, f'image width  {xs[2]} is not multiple of output width  {h}'
    stride1, stride2 = xs[2] // h, xs[3] // w
    assert stride1 == stride2, f'stride of width {stride2} and height {stride1} is unsame'
    stride = stride1
    assert stride in yolov5_anchors, f'stride {stride} should be one of {list(yolov5_anchors.keys())}'
    return yolov5_anchors[stride]

2.2 转换bounding box并进行NMS

def yolov5cut_detect(preds:list, conf_thresh=0.25, iou_thresh=0.45, xs=(1, 3, 640, 640), **kwargs):
    '''裁剪过的yolov5模型的后处理

    - preds是列表实例, 依次为:
      - 尺度1输出  (b, 3 * (4 + 1 + nc), h, w)
      - 尺度2输出  (b, 3 * (4 + 1 + nc), h, w)
      - 尺度3输出  (b, 3 * (4 + 1 + nc), h, w)
    - conf_thresh: 置信度阈值，低于阈值的目标框被忽略
    - iou_thresh: IOU阈值，IOU大于阈值的目标框被忽略，用于NMS
    - xs: 输入尺寸

    yolov5/8默认都是3尺度, 根据需要可以改为2/4尺度, 相应的output数量会相应变化
    输出已经经过sigmoid算子, 可以直接使用
    '''
    check_args(conf_thresh, iou_thresh)
    assert len(preds) > 0, f'Invalid pred length {len(preds)}, valid values should larger than 0'
    # 获取batch数（嵌入式平台中batch一般是1，在PC上的推导时batch可能大于1）
    bs = np.array([pred.shape[0] for pred in preds], dtype=int)
    assert len(np.unique(bs)) == 1, f'Invalid pred batch size {bs}'
    bs = np.unique(bs)[0]

    pack, remain = 1, 0
    assert (len(preds)-remain) % pack == 0, f'Invalid pred length {len(preds)}, valid values % {pack} should be {remain}'
    preds = [ preds[idx:idx+pack] for idx in range(0, len(preds)-remain, pack) ]

    outputs = []
    # 对每个batch单独处理
    for b in range(bs):
     	# 取出batch的数据
    	p = [ p[bs] for p in preds ]
    	# 将特征图转换为bouding box（用置信度阈值进行过滤）
        p = [ yolov5cut_detect_one(o, conf_thresh, yolov5cut_get_anchor(o, xs), xs)
                for o in p ]
        # 合并多个尺度的bouding box
        p = np.concatenate(p, axis=0)
        # NMS
        p = nms(p, iou_thresh)
        outputs.append(p)
    # 0123: 目标框, 4: 置信度，5：类别索引
    return [ (p[:, :4], p[:, 4], p[:, 5], None, None, None) for p in outputs ]

2.3 转换bouding box

def yolov5cut_detect_one(o, conf_thresh, anchor, xs):
    '''转换特征图

	- o: 特征图 （3 * (4 + 1 + nc), h, w)
	  - 01: xy, 以grid中心为基准
	  - 23: wh, 用预选框尺寸和输入图像尺寸进行归一化
	  - 5: 目标置信度, 当前grid是否存在目标
	  - 6:: 类别置信度，nc个类型对应nc个置信度，最终置信度为类别置信度x目标置信度
	- conf_thresh: 目标置信度阈值，忽略低于阈值的目标框
	- anchor: 预选框，每个尺度对应3个预选框，也就是每个grid对3个预选框计算bounding box
	- xs: 输入图像尺寸，用于计算归一化的目标框尺寸
    '''
    c, h, w = o.shape

    # 生成下标矩阵，用于记录被筛选出来的目标框
    # yolov8的目标框是相对grid中心点的偏移，下标矩阵用于记录中心点位置、
    # 结合计算出来的目标框得到真正的目标框
    # 这里生成grid数组是为了用利用numpy的向量计算功能
    # python解释执行循环速度太慢，C/C++不需要生成grid数组
    sx, sy = np.arange(w), np.arange(h)
    sx, sy = np.meshgrid(sx, sy)
    g = np.stack((sy, sx), -1)
    grid = np.zeros(g.shape)
    grid[:] = g
    grid = grid.reshape((h*w, -1))

    # reshape和transpose修改排列方式为(3, bs, h*w, (4 + 1 +nc)), 方便后续numpy处理
    o = o.reshape(c, -1)
    o = o.reshape(3, c//3, -1)
    o = o.transpose((0, 2, 1))
    inter = []
    # 循环处理每个anchor
    for idx, a in enumerate(o):
        ax, ay = anchor[idx*2:(idx+1)*2]

        bbox = a[..., 0:4] # bounding box: xywh
        conf = a[..., 4] # 目标置信度
        cls  = a[..., 5:] # 类别置信度
        conf = conf * cls.max(-1) # 类别置信度(x目标置信度)
        cls = cls.argmax(-1) # 类别索引

        # 筛选目标
        s = conf > conf_thresh
        bbox, conf, cls, grid = bbox[s], conf[s], cls[s], grid[s]
        output = np.zeros((len(bbox), 6))
        for idx, (_bbox, _conf, _cls, (row, col)) in enumerate(zip(bbox, conf, cls, grid)):
            # bounding box的回归值要经过以下计算得到归一化xyxy
            d1, d2, d3, d4 = _bbox * 2
            cx = (d1 - 0.5 + col) / w
            cy = (d2 - 0.5 + row) / h
            ow = d3 * d3 * ax / xs[-1]
            oh = d4 * d4 * ay / xs[-2]
            x1, y1 = cx - ow / 2, cy - oh / 2
            x2, y2 = cx + ow / 2, cy + oh / 2
            output[idx] = (x1, y1, x2, y2, _conf, _cls)
        inter.append(output)
    return np.concatenate(inter, axis=0)

模型转换

在嵌入式处理器上运行YOLOv5模型一般需要将ONNX格式转换为处理器自有的格式，主要解决以下问题：

量化
方便加载模型文件时获取模型推导时需要的计算单元和存储资源（Cache和DDR）

量化问题一般通过在转换时输入典型图片推导计算中间节点和输出节点的输出，以便得到合适的量化位数和scale系数。

通过量化将浮点数转化为定点数通常会带来精度损失，甚至导致不收敛，表现到YOLOv5模型推导中，可能导致目标框和置信度误差。

除量化之外，需要注意的主要有：

输入图片尺寸（转换工具会在模型输入之前增加缩放功能）
输入格式（YOLOv5的输入是RGB格式，但是采集图像经过ISP处理之后一般是YUV420 NV12，转换工具可能会在模型输入之前增加色彩空间转换功能，否则需要自行调用相关硬件模块执行色彩空间转换）
输入scale（YOLOv5的输入需要进行0-1归一化，因此scale系数是1/255）
DC sub（YOLOv5输入是0-1归一化，因此DC sub是0）

模型转换一般通过处理器厂商提供的转换工具完成，参考工具手册即可。这部分工作受保密协议的限制不便公开，技术问题可以私信交流。

C/C++调用代码

C/C++调用代码在嵌入式平台上运行，与python代码的逻辑一一对应，不同的地方有：

模型加载
1. 平台相关的初始化：初始化加速器、设置加速器资源调度模式、初始化内存、初始化辅助硬件模块如色彩空间转换、resize等。
2. 加载模型文件：利用平台相关代码（一般是处理器厂商提供的AI代码库）加载转换后的模型文件，读取模型中的参数信息，包括：输入通道数和维度信息、输出通道数和维度信息、工作缓存大小等。
图片加载
1. 嵌入式处理器一般都集成了JPEG编解码器，如果调用JPEG硬件解码比较方便，可以在模型调试程序中调用JPEG解码器以便可以直接加载JPEG文件；替代方案是移植Opencv或者ffmepg库实现软解码；对于调试程序来说以上两种方案都有些复杂，可以在PC上使用ffmpeg将jpeg文件转换为YUV格式
2. yuv数据要加载到物理内存才能被NPU加速器使用，一般使用厂商提供的AI代码库分配物理内存

ffmpeg -y -i 1280x720-1002-zidane.jpg -s 1280x720 -pix_fmt nv12 1280x720-nv12-1002.yuv

运行模型：使用AI库的接口运行模型并获取结果（一段或几段物理内存，可能带cache，物理内存的格式可以在模型运行完之后获取或者在加载模型时获取，一般来说输出缓存的定义只与模型本身有关）
后处理：与python调用代码不同的地方是，C/C++调用代码需要自行处理量化相关的问题，一般有量化位数和scale两个系数。

转换代码

转换目标框

/**
 * @brief 转换一个bounding box
 * @note 模型输出的缓存格式
 *     - 预选框1、2、3顺序排列
 *     - 每种数据h * w个数据，按bounding box(xywh)、目标置信度(o)、类别置信度(c1-cn)顺序排列
 *     - [xx...xyy...yww...whh...h][oo...o][c1c1...c1c2c2...c2...cncn...cn]
 * @param box [out] 目标框(左上角&右下角，归一化坐标)
 * @param channel_result [in] 模型输出缓存首地址
 * @param idx [in] GRID索引
 * @param col [in] GRID列索引
 * @param row [in] GRID行索引
 * @param block_size [in] h * w
 *
 * this->_ow, this->_oh: 输出尺寸分别对应w和h
 * this->_iw, this->_ih: 输入尺寸
 * this->_scale: 归一化/量化参数
 */
void yolov5cut_convertor::convert_box(
    common::xyxy& box,
    const int16_t* channel_result, const float* anchor,
    uint32_t idx, uint32_t col, uint32_t row, uint32_t block_size)
{
    float d1 = channel_result[idx + 0 * block_size] * this->_scale * 2;
    float d2 = channel_result[idx + 1 * block_size] * this->_scale * 2;
    float d3 = channel_result[idx + 2 * block_size] * this->_scale * 2;
    float d4 = channel_result[idx + 3 * block_size] * this->_scale * 2;

    float x, y, w, h;
    // 计算bouding box，计算方法与python代码一一对应
    x = (d1 - 0.5 + col)    / this->_ow; // bounding box的左上角，输出结果的x相对于GRID中心位置，用输出尺寸归一化
    y = (d2 - 0.5 + row)    / this->_oh;
    w = d3 * d3 * anchor[0] / this->_iw; // 用预选框计算bounding box尺寸，并用输入尺寸将尺寸归一化
    h = d4 * d4 * anchor[1] / this->_ih;

    // 移到左上角
    x = x - w / 2;
    y = y - h / 2;

    // 防止超出边界
    if (x > 1) x = 1;
    if (x < 0) x = 0;
    if (y > 1) y = 1;
    if (y < 0) y = 0;

    if (x + w > 1) w = 1 - x;
    if (y + h > 1) h = 1 - y;

    box = common::xyxy{x, y, x+w, y+h};
}

整体转换过程与python代码相同，python代码利用numpy的向量计算能力加快速度并简化代码，C++代码需要按照batch、anchor、height、weight四层循环逐层计算
1. 计算目标置信度：偏移地址为4 * h * w，乘以量化系数，如果小于阈值则忽略
2. 计算bounding box
3. 查找并计算分类置信度，并记录类别索引
4. 目标置信度x分类置信度得到置信度，如果置信度小于阈值则忽略
5. 保存bounding box
按照分类对bounding box进行NMS

以上是YOLOv5算法移植的全部流程。

下节预告：YOLOv8相比YOLOv5，模型结构有一些改进，使用anchor-free模型解决了YOLOv5 anchor-base的问题，并且支持分割(segment)和关键点(keypoints)两种任务。YOLOv8的移植流程与YOLOv5大致相同，后续文章会描述YOLOv8移植中与YOLOv5不一样的地方。