游乐游手机版
首页/AI热点日报/热点详情

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

类型:热点整理2026-07-16
在目标检测领域,YOLOv5 始终是备受关注的焦点。自 Ultralytics 于 2020 年 6 月发布以来,凭借在 YOLOv3 基础上的多项改进,YOLOv5 迅速成为工业部署的热门选择。目前 YOLOv5 提供 s、m、l、x 四种版本,相比 YOLOv4,它在保持检测平均精度几乎不变的前

在目标检测领域,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.txtvalid.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_detectquant_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 开发包供客户进行二次开发。

来源:https://m.elecfans.com/article/2136598.html

相关热点

继续查看同栏目近期热点。

延伸阅读

补充最近整理过的热点入口。