YOLOv5训练部署从入门到精通全链路完整实战教程

在目标检测领域，YOLOv5 始终是备受关注的焦点。自 Ultralytics 于 2020 年 6 月发布以来，凭借在 YOLOv3 基础上的多项改进，YOLOv5 迅速成为工业部署的热门选择。目前 YOLOv5 提供 s、m、l、x 四种版本，相比 YOLOv4，它在保持检测平均精度几乎不变的前提下，拥有更小的权重文件，训练和推理速度也更快。整个网络架构由输入端、Backbone、Neck 和 Head 四个部分组成。

本教程将完整演示从 YOLOv5 模型训练到部署在 EASY-EAI-Nano（RV1126）开发板上的全过程。关于数据标注的具体方法，可参考我们之前的文章《Labelimg 的安装与使用》。

整体流程如下：

2. 准备数据集

2.1 数据集下载

此处以口罩检测数据集为例。解压完成后将获得以下三个文件：

2.2 生成路径列表

在数据集目录下运行脚本 list_dataset_file.py：

python list_dataset_file.py

执行结果如下图所示：

运行后生成训练样本列表文件 train.txt 和验证样本列表文件 valid.txt：

3. 下载 YOLOv5 训练源码

在 PC 端通过 Git 工具克隆远程仓库：

克隆完成后，目录结构如下：

4. 训练算法模型

切换至 YOLOv5 工作目录，以训练口罩检测模型为例。需要修改 data/mask.yaml 中的 train.txt 和 valid.txt 路径：

然后执行训练脚本：

python train.py --data mask.yaml --cfg yolov5s.yaml --weights "" --batch-size 64

开始训练，效果如下图所示：

算法的精度结果可查看 ./runs/train/results.csv。

5. 在 PC 端进行模型预测

训练完成后，./runs/train/exp/weights/best.pt 即为验证集上效果最佳的模型。可以使用以下命令进行初步预测：

python detect.py --source data/images --weights ./runs/train/exp/weights/best.pt --conf 0.5

6. pt 模型转换为 onnx 模型

要将模型部署到 EASY-EAI-Nano，需要先转换为 RKNN 模型。而在转换为 RKNN 之前，通常先转为 ONNX 格式，同时生成 best.anchors.txt：

python export.py --include onnx --rknpu RV1126 --weights ./runs/train/exp/weights/best.pt

生成的文件如下图所示：

7. 转换为 RKNN 模型环境搭建

ONNX 模型需要进一步转换为 RKNN 格式才能在 EASY-EAI-Nano 上运行。因此，首先需要搭建 rknn-toolkit 模型转换工具的环境。当然，TensorFlow、TensorFlow Lite、Caffe、Darknet 等框架的模型转换方法类似，此处以 ONNX 为例进行说明。

7.1 概述

模型转换环境的搭建流程如下：

7.2 下载模型转换工具

为确保工具顺利运行，请下载 rknn-toolkit-v1.7.1 的 Docker 镜像（文件名：rknn-toolkit-1.7.1-docker.tar.gz）。

7.3 将工具移至 Ubuntu 18.04

将下载好的 Docker 镜像移动到虚拟机的 rknn-toolkit 目录下：

7.4 运行模型转换工具环境

7.4.1 打开终端

在该目录下打开终端：

7.4.2 加载 Docker 镜像

docker load --input /home/developer/rknn-toolkit/rknn-toolkit-1.7.1-docker.tar.gz

7.4.3 进入镜像 bash 环境

docker run -t -i --privileged -v /dev/bus/usb:/dev/bus/usb rknn-toolkit:1.7.1 /bin/bash

进入后如下图所示：

7.4.4 测试环境

输入 python 加载相关库，尝试导入 rknn 库，若出现如下提示则表示环境测试成功：

至此，模型转换工具环境搭建完毕。

8. RKNN 模型转换流程介绍

EASY EAI Nano 支持以 .rknn 为后缀的模型进行评测和运行。常见的 TensorFlow、TensorFlow Lite、Caffe、Darknet、ONNX 和 PyTorch 模型均可通过官方 toolkit 工具转换为 RKNN 模型。其他框架训练的模型，也可以先转成 ONNX 再转 RKNN。转换流程如下图所示：

8.1 模型转换 Demo 下载

将 model_convert.tar.bz2 解压到虚拟机中：

8.2 进入模型转换工具 Docker 环境

执行以下命令，将工作目录映射到 Docker 镜像内（其中 /home/developer/rknn-toolkit/model_convert 为工作区域，/test 为 Docker 内的映射路径，/dev/bus/usb 用于映射 USB 设备）：

docker run -t -i --privileged -v /dev/bus/usb:/dev/bus/usb -v /home/developer/rknn-toolkit/model_convert:/test rknn-toolkit:1.7.1 /bin/bash

执行成功后如上图所示：

8.3 模型转换操作说明

8.3.1 模型转换 Demo 目录结构

Demo 由 mask_object_detect 和 quant_dataset 两部分组成。coco_object_detect 存放软件脚本，quant_dataset 存放量化所需的数据：

mask_object_detect 文件夹内容如下：

8.3.2 生成量化图片列表

在 Docker 环境中切换至模型转换工作目录：

cd /test/mask_object_detect/

执行 gen_list.py 生成量化图片列表：

python gen_list.py

生成的“量化图片列表”如下：

8.3.3 ONNX 模型转换为 RKNN 模型

rknn_convert.py 脚本默认进行 int8 量化，代码如下：

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN

ONNX_MODEL = 'best.onnx'
RKNN_MODEL = './yolov5_mask_rv1126.rknn'
DATASET = './pic_path.txt'
QUANTIZE_ON = True

if __name__ == '__main__':
    rknn = RKNN(verbose=True)
    if not os.path.exists(ONNX_MODEL):
        print('model not exist')
        exit(-1)
    print('--> Config model')
    rknn.config(reorder_channel='0 1 2',
                mean_values=[[0, 0, 0]],
                std_values=[[255, 255, 255]],
                optimization_level=3,
                target_platform='rv1126',
                output_optimize=1,
                quantize_input_node=QUANTIZE_ON)
    print('done')
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print('Load yolov5 failed!')
        exit(ret)
    print('done')
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
    if ret != 0:
        print('Build yolov5 failed!')
        exit(ret)
    print('done')
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export yolov5rknn failed!')
        exit(ret)
    print('done')

将 ONNX 模型 best.onnx 放入 mask_object_detect 目录，然后执行转换脚本：

python rknn_convert.py

生成的模型如下图所示，该模型可在 RKNN 环境以及 EASY EAI Nano 上运行：

8.3.4 运行 RKNN 模型

使用 yolov5_mask_test.py 脚本可在 PC 端的 Docker 环境内运行 RKNN 模型：

脚本代码如下：

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
import random
from rknn.api import RKNN

RKNN_MODEL = 'yolov5_mask_rv1126.rknn'
IMG_PATH = './test.jpg'
DATASET = './dataset.txt'
BOX_THRESH = 0.25
NMS_THRESH = 0.6
IMG_SIZE = 640
CLASSES = ("head", "mask")

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def xywh2xyxy(x):
    y = np.copy(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2
    y[:, 1] = x[:, 1] - x[:, 3] / 2
    y[:, 2] = x[:, 0] + x[:, 2] / 2
    y[:, 3] = x[:, 1] + x[:, 3] / 2
    return y

def process(input, mask, anchors):
    anchors = [anchors[i] for i in mask]
    grid_h, grid_w = map(int, input.shape[0:2])
    box_confidence = sigmoid(input[..., 4])
    box_confidence = np.expand_dims(box_confidence, axis=-1)
    box_class_probs = sigmoid(input[..., 5:])
    box_xy = sigmoid(input[..., :2]) * 2 - 0.5
    col = np.tile(np.arange(0, grid_w), grid_w).reshape(-1, grid_w)
    row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_h)
    col = col.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
    row = row.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
    grid = np.concatenate((col, row), axis=-1)
    box_xy += grid
    box_xy *= int(IMG_SIZE / grid_h)
    box_wh = pow(sigmoid(input[..., 2:4]) * 2, 2)
    box_wh = box_wh * anchors
    box = np.concatenate((box_xy, box_wh), axis=-1)
    return box, box_confidence, box_class_probs

def filter_boxes(boxes, box_confidences, box_class_probs):
    box_scores = box_confidences * box_class_probs
    box_classes = np.argmax(box_class_probs, axis=-1)
    box_class_scores = np.max(box_scores, axis=-1)
    pos = np.where(box_confidences[..., 0] >= BOX_THRESH)
    boxes = boxes[pos]
    classes = box_classes[pos]
    scores = box_class_scores[pos]
    return boxes, classes, scores

def nms_boxes(boxes, scores):
    x = boxes[:, 0]
    y = boxes[:, 1]
    w = boxes[:, 2] - boxes[:, 0]
    h = boxes[:, 3] - boxes[:, 1]
    areas = w * h
    order = scores.argsort()[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x[i], x[order[1:]])
        yy1 = np.maximum(y[i], y[order[1:]])
        xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
        yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])
        w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)
        h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
        inter = w1 * h1
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(ovr <= NMS_THRESH)[0]
        order = order[inds + 1]
    keep = np.array(keep)
    return keep

def yolov5_post_process(input_data):
    masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
    anchors = [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45],
               [59, 119], [116, 90], [156, 198], [373, 326]]
    boxes, classes, scores = [], [], []
    for input, mask in zip(input_data, masks):
        b, c, s = process(input, mask, anchors)
        b, c, s = filter_boxes(b, c, s)
        boxes.append(b)
        classes.append(c)
        scores.append(s)
    boxes = np.concatenate(boxes)
    boxes = xywh2xyxy(boxes)
    classes = np.concatenate(classes)
    scores = np.concatenate(scores)
    nboxes, nclasses, nscores = [], [], []
    for c in set(classes):
        inds = np.where(classes == c)
        b = boxes[inds]
        c = classes[inds]
        s = scores[inds]
        keep = nms_boxes(b, s)
        nboxes.append(b[keep])
        nclasses.append(c[keep])
        nscores.append(s[keep])
    if not nclasses and not nscores:
        return None, None, None
    boxes = np.concatenate(nboxes)
    classes = np.concatenate(nclasses)
    scores = np.concatenate(nscores)
    return boxes, classes, scores

def scale_coords(x1, y1, x2, y2, dst_width, dst_height):
    dst_top, dst_left, dst_right, dst_bottom = 0, 0, 0, 0
    gain = 0
    if dst_width > dst_height:
        image_max_len = dst_width
        gain = IMG_SIZE / image_max_len
        resized_height = dst_height * gain
        height_pading = (IMG_SIZE - resized_height) / 2
        y1 = (y1 - height_pading)
        y2 = (y2 - height_pading)
    dst_x1 = int(x1 / gain)
    dst_y1 = int(y1 / gain)
    dst_x2 = int(x2 / gain)
    dst_y2 = int(y2 / gain)
    return dst_x1, dst_y1, dst_x2, dst_y2

def plot_one_box(x, img, color=None, label=None, line_thickness=None):
    tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)
        cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)

def draw(image, boxes, scores, classes):
    for box, score, cl in zip(boxes, scores, classes):
        x1, y1, x2, y2 = box
        print('class: {}, score: {}'.format(CLASSES[cl], score))
        print('box coordinate x1,y1,x2,y2: [{}, {}, {}, {}]'.format(x1, y1, x2, y2))
        x1 = int(x1)
        y1 = int(y1)
        x2 = int(x2)
        y2 = int(y2)
        dst_x1, dst_y1, dst_x2, dst_y2 = scale_coords(x1, y1, x2, y2, image.shape[1], image.shape[0])
        plot_one_box((dst_x1, dst_y1, dst_x2, dst_y2), image, label='{0} {1:.2f}'.format(CLASSES[cl], score))

def letterbox(im, new_shape=(640, 640), color=(0, 0, 0)):
    shape = im.shape[:2]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    ratio = r, r
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]
    dw /= 2
    dh /= 2
    if shape[::-1] != new_unpad:
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
    return im, ratio, (dw, dh)

if __name__ == '__main__':
    rknn = RKNN(verbose=True)
    print('--> Loading model')
    ret = rknn.load_rknn(RKNN_MODEL)
    if ret != 0:
        print('load rknn model failed')
        exit(ret)
    print('done')
    print('--> Init runtime environment')
    ret = rknn.init_runtime()
    if ret != 0:
        print('Init runtime environment failed')
        exit(ret)
    print('done')
    img = cv2.imread(IMG_PATH)
    letter_img, ratio, (dw, dh) = letterbox(img, new_shape=(IMG_SIZE, IMG_SIZE))
    letter_img = cv2.cvtColor(letter_img, cv2.COLOR_BGR2RGB)
    print('--> Running model')
    outputs = rknn.inference(inputs=[letter_img])
    input0_data = outputs[0]
    input1_data = outputs[1]
    input2_data = outputs[2]
    input0_data = input0_data.reshape([3, -1] + list(input0_data.shape[-2:]))
    input1_data = input1_data.reshape([3, -1] + list(input1_data.shape[-2:]))
    input2_data = input2_data.reshape([3, -1] + list(input2_data.shape[-2:]))
    input_data = list()
    input_data.append(np.transpose(input0_data, (2, 3, 0, 1)))
    input_data.append(np.transpose(input1_data, (2, 3, 0, 1)))
    input_data.append(np.transpose(input2_data, (2, 3, 0, 1)))
    boxes, classes, scores = yolov5_post_process(input_data)
    if boxes is not None:
        draw(img, boxes, scores, classes)
    cv2.imwrite('./result.jpg', img)
    rknn.release()

执行后得到 result.jpg，检测效果如下：

8.3.5 模型预编译

RKNN 模型直接使用 NPU API 在 EASY EAI Nano 上加载时，启动速度较慢。在评估模型精度无误后，建议进行一次预编译。预编译需要借助 EASY EAI Nano 主板的环境，因此务必确保开发板通过 adb 与 Ubuntu 稳定连接。

板子端的接线如下图所示，拨码开关需拨至 adb 模式：

虚拟机需要能够识别 adb 设备：

由于 Ubuntu 环境与 Docker 环境对 adb 设备资源存在竞争，需先关闭 Ubuntu 环境下的 adb 服务，并在 Docker 中安装 adb 软件包。以下指令分别在 Ubuntu 环境和 Docker 环境中执行：

在 Docker 环境中执行 adb devices，出现如下提示则表示设备连接成功：

运行 precompile_rknn.py 脚本进行预编译：

python precompile_rknn.py

执行效果如下图所示，生成预编译模型 yolov5_mask_rv1126_pre.rknn：

至此，预编译部署完成。最终生成的 int8 量化模型如下图所示：

9. 模型部署示例

9.1 模型部署示例介绍

本节展示 YOLOv5 模型在 EASY EAI Nano 上的实际部署过程。注意，该模型仅经过简单训练，仅供示例参考，不保证精度。

9.2 准备工作

9.2.1 硬件准备

需要 EASY EAI Nano 开发板、microUSB 数据线以及一台搭载 Linux 操作系统的电脑。确保开发板与 Linux 系统之间已建立 adb 连接。

9.2.2 交叉编译环境准备

本示例需要交叉编译环境，可参考在线文档“入门指南 / 开发环境准备 / 安装交叉编译工具链”。

9.2.3 文件下载

下载 YOLOv5 C Demo 示例文件，解压后如下图所示：

9.3 在 EASY EAI Nano 运行 YOLOv5 demo

9.3.1 解压 YOLOv5 demo

将下载的程序包移至 Ubuntu 环境后，执行以下命令解压：

tar -xvf yolov5_detect_C_demo.tar.bz2

9.3.2 编译 YOLOv5 demo

./build.sh

编译成功后的现象：

9.3.3 执行 YOLOv5 demo

将可执行程序推送至开发板端：

adb push yolov5_detect_demo_release/ /userdata

登录到开发板并执行：

adb shell
cd /userdata/yolov5_detect_demo_release/
./yolov5_detect_demo

执行结果如下图所示，算法单帧处理时间约为 50ms：

取回测试图片：

adb pull /userdata/yolov5_detect_demo_release/result.jpg .

测试结果如下：

10. 基于摄像头的 AI Demo

10.1 摄像头 Demo 介绍

本节展示 YOLOv5 模型在 EASY EAI Nano 上运行摄像头实时识别的过程。同样，该模型仅供示例使用，不保证精度。

10.2 准备工作

10.2.1 硬件准备

需要 EASY-EAI-Nano 人工智能开发套件（包含开发板、双目摄像头、5 寸高清屏幕、microUSB 数据线），以及搭载 Linux 操作系统的电脑。确保 adb 连接正常。

10.2.2 交叉编译环境准备

本示例需要交叉编译环境，可参考在线文档“入门指南 / 开发环境准备 / 安装交叉编译工具链”。

10.2.3 文件下载

摄像头识别 Demo 的程序源码下载后解压如下图所示：

10.3 在 EASY EAI Nano 运行 YOLOv5 camera demo

10.3.1 解压 YOLOv5 camera demo

tar -xvf yolov5_detect_camera_demo.tar.bz2

10.3.2 编译 YOLOv5 camera demo

./build.sh

编译成功后的现象：

10.3.3 执行 YOLOv5 camera demo

将可执行程序推送至开发板端：

adb push yolov5_detect_camera_demo_release/ /userdata

登录到开发板并执行：

adb shell
cd /userdata/yolov5_detect_camera_demo_release/
./yolov5_detect_camera_demo

测试结果如下图所示：

硬件使用

本教程使用的开发板为 EASY EAI Nano（RV1126）。

EASY EAI Nano 基于瑞芯微 RV1126 处理器设计，具备四核 CPU @1.5GHz 与 NPU @2Tops 的 AI 边缘计算能力。实现 AI 运算的功耗仅为 GPU 方案的不到 10%。配套的 AI 算法工具非常完善，支持 TensorFlow、PyTorch、Caffe、MxNet、DarkNet、ONNX 等主流框架的直接转换与部署。软硬件开发资料丰富，外设接口齐全，并提供多种功能配件可选。集成有以太网、Wi-Fi 等通信外设，摄像头、显示屏（带电容触摸）、喇叭、麦克风等交互外设，2 路 USB Host、1 路 USB Device 调试接口，以及协议串口、TF 卡、IO 拓展接口（兼容树莓派 / Jetson Nano 拓展接口）等通用外设。内置人脸识别、安全帽监测、人体骨骼点识别、火焰检测、车辆检测等多种 AI 算法，并提供完整的 Linux 开发包供客户进行二次开发。

你是一名 AI 行业编辑，请围绕下面这条热点输出一份资讯解读：
热点：YOLOv5训练部署从入门到精通全链路完整实战教程要求：
1. 先用一句话解释这条热点在讲什么
2. 再总结它为什么重要
3. 说明会影响哪些 AI 产品或内容方向
4. 最后给出 3 个适合资讯站使用的标题