今天,我将展示如何使用令人印象深刻的Hailo AI Hat在树莓派5上训练、编译和部署自定义模型。注意:文章内的链接可能需要科学上网。
Hailo AI Hat
根据你的设置,在树莓派5的CPU上运行YOLO每秒可以提供1.5到8帧(FPS)。尽管对于这样一个小设备来说,这一性能已经相当出色,但对于许多实时应用来说,这还不够快。如果你需要更高的性能,就需要外部硬件。目前,Hailo为树莓派5设计的AI Hat是一个极佳的选择。
Hailo是一家专注于开发人工智能硬件的芯片制造商。在这个故事中,我感兴趣的是他们专门为树莓派5打造的AI Hat+设备:
我的树莓派 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
准备数据
我们最终得到以下结构:
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 .
构建镜像
现在,运行容器:
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
使用交互模式运行容器
这个容器没有nano编辑器。除非你是Vim用户,否则我建议按照以下方式安装nano:
sudoapt updatesudo apt install nano -y
安装完nano后,我们可以继续并设置我们的训练。将datasets文件夹复制到workspace文件夹中:
cp-r /home/hailo/shared/datasets ..
现在,编写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。
创建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将开始全力运转:
我的RTX 4070在燃烧!
保持温度在81°C以下,你就没事。
在我的情况下,这次训练大约用了40分钟。
作为参考,使用另一台配备GTX 1080的机器大约需要2小时。
最后,你会得到类似这样的结果:
训练结束
这意味着训练已经完成。我们可以在runs/exp0文件夹中检查结果。将这个文件夹复制到共享区域:
mkdir/home/hailo/shared/runscp-r runs/exp0 /home/hailo/shared/runs/
你最终会得到一个这样的文件夹:
训练结果文件夹
我们可以检查训练结果:
训练结果
比较第一行图表(训练性能)和第二行图表(验证性能),我们发现模型没有过拟合。
值得一提的是,使用不同的验证实例集是评估模型质量的首要要求。
可以使用常规的机器学习工程技术来改进模型,以达到更高的性能。然而,这并不是我们现在的重点。
记住:我们的重点是学习如何在树莓派/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已经生成:
将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
这个命令会给我们一个容器机器内的命令提示符:
运行新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
Hailo 网络下载页面
将它们保存在共享文件夹的某个地方,并将它们复制到容器中:
cp/home/hailo/shared/libs/* .
在安装软件之前,为Python创建一个虚拟环境并激活它:
python -m venv .venvsource.venv/bin/activate
然后,安装Hailo RT包:
dpkg-i ./hailort_4.21.0_amd64.deb
安装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
安装Hailo内容
坚持住!最困难的部分现在来了:将best.onnx文件编译成best.hef文件。
要使这工作,我们需要将hailo_model_zoo/cfg/postprocess_config/yolov5s_nms_config.json中的类别数更改为2:
请注意,在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优化并编译你的模型:
用hailomz编译模型
如果一切顺利,你会得到以下消息:
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" ]}
使用这个视频的结果是:
即使在高清分辨率下,对象检测也能达到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