完整指南:如何使用树莓派5、Hailo AI Hat、YOLO、Docker进行自定义数据集训练?

    科创经济 朗峰江湖 2025-06-30 4296 次浏览

    今天,我将展示如何使用令人印象深刻的Hailo AI Hat在树莓派5上训练、编译和部署自定义模型。注意:文章内的链接可能需要科学上网。

    Hailo AI Hat

    根据你的设置,在树莓派5的CPU上运行YOLO每秒可以提供1.5到8帧(FPS)。尽管对于这样一个小设备来说,这一性能已经相当出色,但对于许多实时应用来说,这还不够快。如果你需要更高的性能,就需要外部硬件。目前,Hailo为树莓派5设计的AI Hat是一个极佳的选择。

    Hailo是一家专注于开发人工智能硬件的芯片制造商。在这个故事中,我感兴趣的是他们专门为树莓派5打造的AI Hat+设备:

    24e5df38-53b6-11f0-986f-92fbcf53809c.jpg

    我的树莓派 5带AI Hat+ Hailo8,还有Camera Noir v2

    这款AI Hat有两个版本,一个搭载Hailo-8芯片,据说可以提供26 TOPS(每秒万亿次运算),另一个搭载Hailo-8L芯片,提供13 TOPS。

    在这个故事中,我将使用搭载Hailo8架构的AI Hat版本。

    为什么选择Docker?

    准备好下载并安装千兆字节的第三方库吧。我们将使用Docker容器作为隔离环境,来配置和安装所需的一切,而无需修改主机。

    也就是说,我们需要设置两个不同的Docker容器:

    YOLOv5容器:这个容器有两个任务。首先,我们将使用自定义数据集训练模型。其次,我们将模型转换为ONNX格式。

    Hailo容器:这个容器用于将ONNX文件转换为Hailo的HEF格式。

    尝试使用同一个Docker容器来完成这两项任务,会因为库(如numpy)的冲突而带来不必要的麻烦。相信我,使用Docker来做它设计的事情,可以节省你的时间!

    第一个Dockerfile由Hailo提供。第二个Dockerfile将在这个故事的后面部分提供。

    本例中使用的数据集

    我将使用Tech Zizou标记口罩数据集。你可以在这里找到它。

    nature中发现的Tech Zizou标签口罩数据

    从Kaggle下载文件,并按照以下方式解压:

    mkdirsourceunzip -qq archive.zip -dsource/

    source/obj中的文件结构不符合YOLO的预期。希望下面的代码能解决这个问题:

    importos, shutil, random# preparing the folder structurefull_data_path ='source/obj/'extension_allowed ='.jpg'split_percentage =90images_path ='datasets/images/'ifos.path.exists(images_path): shutil.rmtree(images_path)os.mkdir(images_path)labels_path ='datasets/labels/'ifos.path.exists(labels_path): shutil.rmtree(labels_path)os.mkdir(labels_path)training_images_path = images_path +'train/'validation_images_path = images_path +'val/'training_labels_path = labels_path +'train/'validation_labels_path = labels_path +'val/'os.mkdir(training_images_path)os.mkdir(validation_images_path)os.mkdir(training_labels_path)os.mkdir(validation_labels_path)files = []ext_len =len(extension_allowed)forr, d, f in os.walk(full_data_path): forfile in f: iffile.endswith(extension_allowed): strip = file[0:len(file) - ext_len] files.append(strip)random.shuffle(files)size =len(files) split =int(split_percentage * size /100)print("copying training data")fori inrange(split): strip = files[i] image_file = strip + extension_allowed src_image = full_data_path + image_file shutil.copy(src_image, training_images_path) annotation_file = strip +'.txt' src_label = full_data_path + annotation_file shutil.copy(src_label, training_labels_path)print("copying validation data")fori inrange(split, size): strip = files[i] image_file = strip + extension_allowed src_image = full_data_path + image_file shutil.copy(src_image, validation_images_path) annotation_file = strip +'.txt' src_label = full_data_path + annotation_file shutil.copy(src_label, validation_labels_path)print("finished")

    这段代码假设数据在source/obj/文件夹中,并将输出数据放入datasets文件夹中。将文件命名为tidy_data.py,并按照以下方式运行:

    mkdirdatasetspython tidy_data.py

    准备数据

    我们最终得到以下结构:

    252ffc4e-53b6-11f0-986f-92fbcf53809c.png

    Yolo预期的文件夹结构

    这里有一些需要注意的事项:

    这个数据集只有两个类别:戴口罩和不戴口罩。

    只有1359张训练图像和151张验证图像。

    训练数据量很小。仅使用这些数据从头开始训练模型将产生非常差的模型,这种情况称为过拟合。

    我们在这里不深入探讨建模细节。无论如何,为了简化事情,我们将使用一种称为迁移学习的技术,即在训练开始前,将预训练的权重输入到模型中。特别是,我们将使用Ultralytics提供的使用COCO数据库训练的权重。

    YOLOv5

    Ultralytics的最新YOLO版本是11。

    它比YOLOv5更快、更准确。但这并不意味着YOLOv5已经过时。实际上,Ultralytics明确表示,在某些特定场景下,YOLOv5是更优的选择。

    在这个故事中,我有充分的理由避免使用YOLO 11:Hailo堆栈目前还不支持YOLO 11。

    如果你真的不想使用YOLOv5,你可以轻松地将这个故事改编为使用YOLO 8。

    Linux,朋友,Linux!

    这个故事使用Linux,具体来说是Ubuntu LTS。

    对于人工智能开发,我推荐使用Ubuntu 20.04或22.04。LTS一路相伴!

    任务简报

    整个过程由三个简单的步骤组成:

    步骤1:训练自定义模型:在这一步中,我们使用自定义数据加上预训练的YOLOv5权重来训练模型,以执行我们的检测任务(在我们的例子中,是检测戴或不戴口罩的脸)。这一步的输出是一个pytorchbest.pt文件。这个文件只包含我们模型的参数值。

    步骤2:将best.pt转换为ONNX格式:ONNX是一种用于机器学习模型的开放格式。这一步的输出是一个best.onnx文件。

    步骤3:将best.onnx转换为HEF格式:Hailo可执行格式是一种专门为在Hailo芯片上运行而高度优化的模型。在这一步中,我们将ONNX文件转换为HEF文件。

    一旦我们有了.hef格式的模型,我们只需将其部署到树莓派上并进行测试。

    步骤1:训练你的自定义数据

    为Hailo架构训练模型并没有什么新奇之处。你可以像往常一样训练你的模型。

    如果你已经有了模型,就跳过这一步。否则,请继续阅读。

    首先,如果系统中还没有安装Docker,请安装它。

    同时,安装NVIDIA Container Toolkit。

    Hailo在GitHub上共享了一个包含所需资源的仓库。克隆它:

    gitclonehttps://github.com/hailo-ai/hailo_model_zoo

    克隆 the hailo_model_zoo 仓库

    我们在hailo_model_zoo/training/yolov5文件夹中寻找YOLOv5训练的Dockerfile。移动到这个文件夹,并使用以下命令构建镜像:

    cdhailo_model_zoo/training/yolov5docker build -t yolov5:v0 .

    256669be-53b6-11f0-986f-92fbcf53809c.png

    构建镜像

    现在,运行容器:

    docker run-it--name custom_training--gpus all--ipc=host-v/home/doleron/hailo/shared:/home/hailo/shared yolov5:v0

    简而言之,-it标志要求Docker以交互模式运行容器,这对于后续执行命令是必要的。

    参数-v /home/doleron/hailo/shared:/home/hailo/shared将我机器上的/home/doleron/hailo/shared文件夹映射到容器机器上的/home/hailo/shared文件夹。

    --gpus all指示Docker使用主机上可用的任何GPU。

    现在,我们在容器内部。我们可以检查/home/hailo/shared的内容,以确保我们的数据集文件在那里:

    ls/home/hailo/shared/ -als

    251bab22-53b6-11f0-986f-92fbcf53809c.png

    使用交互模式运行容器

    这个容器没有nano编辑器。除非你是Vim用户,否则我建议按照以下方式安装nano:

    sudoapt updatesudo apt install nano -y

    安装完nano后,我们可以继续并设置我们的训练。将datasets文件夹复制到workspace文件夹中:

    cp-r /home/hailo/shared/datasets ..

    2588c2ac-53b6-11f0-986f-92fbcf53809c.png

    现在,编写data/dataset.yaml文件:

    nano data/dataset.yaml

    这是data/dataset.yaml的内容:

    train: ../datasets/images/trainval: ../datasets/images/valnc: 2names: 0:'using mask' 1:'without mask'

    按control-x,y,然后Enter保存文件并退出nano。

    259cd616-53b6-11f0-986f-92fbcf53809c.png

    创建data/dataset.yaml文件

    是时候训练我们的模型了!确保你在/workspace/yolov5文件夹中,并输入:

    python train.py--img640--batch16--epochs100--datadataset.yaml--weightsyolov5s.pt

    如果你遇到类似RuntimeError: CUDA out of memory的错误,尝试将--batch 16减少到--batch 8或更少。

    我希望你熟悉基本的机器学习术语:batches、epoch等。你可以按照这份指南调整这些超参数。

    https://docs.ultralytics.com/zh/guides/hyperparameter-tuning/

    如果一切顺利,你的GPU将开始全力运转:

    25b44300-53b6-11f0-986f-92fbcf53809c.png

    我的RTX 4070在燃烧!

    25cd3144-53b6-11f0-986f-92fbcf53809c.png

    保持温度在81°C以下,你就没事。

    在我的情况下,这次训练大约用了40分钟。

    作为参考,使用另一台配备GTX 1080的机器大约需要2小时。

    最后,你会得到类似这样的结果:

    25eae5d6-53b6-11f0-986f-92fbcf53809c.png

    训练结束

    这意味着训练已经完成。我们可以在runs/exp0文件夹中检查结果。将这个文件夹复制到共享区域:

    mkdir/home/hailo/shared/runscp-r runs/exp0 /home/hailo/shared/runs/

    你最终会得到一个这样的文件夹:

    25fbf358-53b6-11f0-986f-92fbcf53809c.png

    训练结果文件夹

    我们可以检查训练结果:

    260edfb8-53b6-11f0-986f-92fbcf53809c.png

    训练结果

    比较第一行图表(训练性能)和第二行图表(验证性能),我们发现模型没有过拟合。

    值得一提的是,使用不同的验证实例集是评估模型质量的首要要求。

    可以使用常规的机器学习工程技术来改进模型,以达到更高的性能。然而,这并不是我们现在的重点。

    记住:我们的重点是学习如何在树莓派/Hailo AI Hat上使用这样的模型。

    让我们进入下一步!

    步骤2:将best.pt文件转换为ONNX

    回到容器中,最好的权重文件是runs/exp0/weights/best.pt。我们可以使用以下命令将其转换为ONNX:

    python3 models/export.py --weights runs/exp0/weights/best.pt --img 640

    注意,best.onnx已经生成:

    2643e118-53b6-11f0-986f-92fbcf53809c.png

    将best.pt转换成best.onnx

    将best.onnx复制到主机机器上:

    cpruns/exp0/weights/best.onnx /home/hailo/shared/

    我们已经完成了这个容器的任务。如果你想退出,就退出吧。

    步骤3:将ONNX转换为HEF

    教程中最简单的部分是使用YOLOv5训练自定义模型并将结果文件转换为ONNX。现在,是时候将ONNX文件编译成专有的Hailo可执行格式(HEF)了。

    在任何地方启动一个新的终端,并编写这个Dockerfile:

    # using a CUDA supported Ubuntu 22.04 image as baseFROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04 AS base_cudaENV DEBIAN_FRONTEND=noninteractiveENV PYTHONDONTWRITEBYTECODE=1ENV PYTHONUNBUFFERED=1RUN apt-get update && \ apt-get install -y \ # see: the Hailo DFC user guide python3.10 \ python3.10-dev \ python3.10-venv \ python3.10-distutils \ python3-pip \ python3-tk \ graphviz \ libgraphviz-dev \ libgl1-mesa-glx \ # utilities python-is-python3 \ build-essential \ sudo \ curl \ git \ nano && \ # clean up rm-rf /var/lib/apt/lists/*# update pipRUN python3 -m pip install --upgrade pip setuptools wheelWORKDIR /workspaceARG user=hailoARG group=hailoARG uid=1000ARG gid=1000RUN groupadd --gid$gid$group&& \ adduser --uid$uid--gid$gid--shell /bin/bash --disabled-password --gecos""$user&& \ chmodu+w /etc/sudoers &&echo"$userALL=(ALL) NOPASSWD: ALL">> /etc/sudoers &&chmod-w /etc/sudoers && \ chown-R$user:$group/workspace

    将其保存为Dockerfile,并使用以下命令构建镜像:

    dockerbuild -t hailo_compiler:v0 .

    一旦镜像构建完成,按照以下方式启动容器:

    docker run-it--name compile_onnx_file--gpus all--ipc=host-v/home/doleron/hailo/shared:/home/hailo/shared hailo_compiler:v0

    这个命令会给我们一个容器机器内的命令提示符:

    265c81e6-53b6-11f0-986f-92fbcf53809c.png

    运行新Docker容器

    这看起来像是似曾相识。我们在上一节中刚刚执行了类似的步骤。那又怎样?

    关键在于:我们正在挂载第二个隔离容器来安装Hailo的东西,而不用担心与其他库的冲突。特别是,我们需要安装三个包:

    Hailort:Hailo运行时平台

    Hailort Wheel:Hailort Python库

    Hailo DFC:Hailo数据流编译器

    FOSS社区习惯了开源生态系统。在这个上下文中,一切都可以从公开可用的仓库中安装。然而,Hailo在人工智能市场这个充满挑战和野性的商业世界中运作。因此,他们的软件还不是开源的。希望Hailo的软件至少是免费的。

    要使用Hailo的东西,我们必须在Hailo Network上创建一个账户.

    访问软件下载页面,并下载三个包:

    hailort_4.21.0_amd64.deb

    hailort-4.21.0-cp310-cp310-linux_x86_64.whl

    hailo_dataflow_compiler-3.31.0-py3-none-linux_x86_64.whl

    266fdbe2-53b6-11f0-986f-92fbcf53809c.png

    Hailo 网络下载页面

    将它们保存在共享文件夹的某个地方,并将它们复制到容器中:

    cp/home/hailo/shared/libs/* .

    在安装软件之前,为Python创建一个虚拟环境并激活它:

    python -m venv .venvsource.venv/bin/activate

    然后,安装Hailo RT包:

    dpkg-i ./hailort_4.21.0_amd64.deb

    2680f2a6-53b6-11f0-986f-92fbcf53809c.png

    安装Hailo RT

    接下来,安装Hailo RT Python API:

    pipinstall ./hailort-4.21.0-cp310-cp310-linux_x86_64.whl

    现在,安装Hailo DFC:

    pipinstall ./hailo_dataflow_compiler-3.31.0-py3-none-linux_x86_64.whl

    注意,包版本表示在这个故事编写时Hailo软件的当前阶段。它们必须与容器Python版本(3.10)匹配。

    我们还没完成。我们必须克隆并安装hailo_model_zoo:

    gitclonehttps://github.com/hailo-ai/hailo_model_zoo.gitcdhailo_model_zoopip install -e .

    检查hailomz是否正确设置:

    hailomz--version

    269080ea-53b6-11f0-986f-92fbcf53809c.png

    安装Hailo内容

    坚持住!最困难的部分现在来了:将best.onnx文件编译成best.hef文件。

    要使这工作,我们需要将hailo_model_zoo/cfg/postprocess_config/yolov5s_nms_config.json中的类别数更改为2:

    26aab8c0-53b6-11f0-986f-92fbcf53809c.png

    请注意,在hailo_model_zoo仓库中有一个hailo_model_zoo文件夹!

    在开始编译器之前,设置USER环境变量:

    exportUSER=hailo

    现在,按照以下方式调用hailomz:

    hailomz compile--ckpt/home/hailo/shared/best.onnx--calib-path/home/hailo/shared/datasets/images/train/--yaml hailo_model_zoo/cfg/networks/yolov5s.yaml

    慢慢来。等待10分钟,让hailomz优化并编译你的模型:

    26c1f51c-53b6-11f0-986f-92fbcf53809c.png

    用hailomz编译模型

    如果一切顺利,你会得到以下消息:

    251bab22-53b6-11f0-986f-92fbcf53809c.png

    HEF 编译

    注意,将ONNX转换为HEF包括一个新元素:校准图像。校准图像是Hailo编译器用于优化模型的特性空间的示例。我在这里没有找到任何文档,但一旦hailomz编译器警告我使用超过1024个实例,使用相同的训练集似乎就能工作。

    将yolov5s.hef复制到共享区域:

    cpyolov5s.hef /home/hailo/shared/

    最困难的部分已经完成。我们可以退出容器实例。

    在树莓派5上部署模型

    将yolov5s.hef复制到树莓派上。

    在树莓派上运行Hailo应用程序的细节超出了这个故事的范围。

    在树莓派上,运行以下命令:

    gitclonehttps://github.com/hailo-ai/hailo-rpi5-examples.gitcdhailo-rpi5-examplessourcesetup_env.shpython basic_pipelines/detection.py --labels-json custom.json --hef-path /home/pi/Documents/yolov5s.hef --input /home/pi/Documents/videoplayback.mp4 -f

    其中custom.json是:

    { "detection_threshold":0.5, "max_boxes":200, "labels":[ "unlabeled", "with mask", "without mask" ]}

    使用这个视频的结果是:

    26ea9206-53b6-11f0-986f-92fbcf53809c.png

    即使在高清分辨率下,对象检测也能达到30 fps。这是令人印象深刻的!您可以探索其他输入类型,例如:

    python basic_pipelines/detection.py --labels-json custom.json --hef-path /home/pi/Documents/yolov5s.hef --input usb -f

    或者

    python basic_pipelines/detection.py --labels-json custom.json --hef-path /home/pi/Documents/yolov5s.hef --input rpi -f

    查看Hailo RPI示例仓库以获取更多参数和用法示例。

    https://github.com/hailo-ai/hailo-rpi5-examples

    使用其他YOLO版本

    值得注意的是,在撰写本文时,Hailo模型编译器仅与YOLO3、YOLO4、YOLOv5、YOLOv8和YOLOX进行了测试。

    查看Hailo开发者专区,了解Hailo编译器何时将支持更早的YOLO版本。

    结论

    我们展示了使用Hailo AI Hat在树莓派5上训练自定义数据集、编译和部署模型的完整步骤序列。

    我期待着弄清楚AI Hat能做什么。但这是另一个故事的话题了。

    原文地址:

    https://pub.towardsai.net/custom-dataset-with-hailo-ai-hat-yolo-raspberry-pi-5-and-docker-0d88ef5eb70f