在 Arduino 上为 microTVM 训练视觉模型
单击 此处 下载完整的示例代码
作者:Gavin Uberti
本教程介绍如何训练 MobileNetV1 模型以适应嵌入式设备,以及如何使用 TVM 将这些模型部署到 Arduino。
推荐用 Jupyter Notebook 查看本教程的代码,可以用本页底部的链接下载并运行,或者使用 Google Colab 免费在线查看。
背景简介
构建物联网设备时,通常想让它们能够看到并理解它们周围的世界。可以采取多种形式,但通常设备也会想知道某种物体是否在其视野中。
例如,安全摄像头可能会寻找人,因此它可以决定是否将视频保存 到内存中。红绿灯可能会寻找汽车,这样它就可以判断哪个信号灯应该首先改变。或者森林相机可能会寻找一种动物,从而估计动物种群的数量。
为使这些设备价格合理,我们希望为这些设备配置一个低成本处理器,如 nRF52840(在 Mouser 上每个售价 5 美元)或 RP2040(每个只需 1.45 美元)。
这些设备的内存非常小(~250 KB RAM),这意味着传统的边缘 AI 视觉模型(如 MobileNet 或 EfficientNet)都不能够运行。本教程将展示如何修改这些模型以解决此问题。然后,使用 TVM 为 Arduino 编译和部署。
安装必要软件
本教程使用 TensorFlow(Google 创建的一个广泛使用的机器学习库)来训练模型。TensorFlow 是一个非常底层的库,因此要用 Keras 接口来从 TensorFlow 获取信息。还会用 TensorFlow Lite 量化模型,因为 TensorFlow 本身不支持这一点。
生成模型后,使用 TVM 对其进行编译和测试。为避免必须从源代码构建,需要安装 tlcpack
- TVM 的社区构建工具,还要安装 imagemagick
和 curl
来预处理数据:
%%bash
pip install -q tensorflow tflite
pip install -q tlcpack-nightly -f https://tlcpack.ai/wheels
apt-get -qq install imagemagick curl
# 为 Nano 33 BLE 安装 Arduino CLI 和库
curl -fsSL https://raw.githubusercontent.com/arduino/arduino-cli/master/install.sh | sh
/content/bin/arduino-cli core update-index
/content/bin/arduino-cli core install arduino:mbed_nano
使用 GPU
本教程演示如何训练神经网络,训练神经网络需要大量的计算能力,使用 GPU 训练速度会更快。若在 Google Colab 上查看本教程,可通过 Runtime->Change runtime type 并选择“GPU”作为硬件加速器来启用 GPU。若在本地运行,则可按照 TensorFlow 指南 进行操作。
使用以下代码测试 GPU 是否安装:
import tensorflow as tf
if not tf.test.gpu_device_name():
print("No GPU was detected!")
print("Model training will take much longer (~30 minutes instead of ~5)")
else:
print("GPU detected - you're good to go.")
输出结果:
No GPU was detected!
Model training will take much longer (~30 minutes instead of ~5)
选择工作目录
选择工作目录,把图像数据集、训练好的模型和最终的 Arduino sketch 都放在此目录下,如果在 Google Colab 上运行,所有内容将保存在/root
(又名 ~
)中,若在本地运行,可将其存储在其他地方。注意,此变量仅影响 Python 脚本 - 还必须调整 Bash 命令。
import os
FOLDER = "/root"
下载数据
卷积神经网络通过大量图像以及标签来学习,为获得这些图像,需要一个公开可用的数据集,这个数据集中要包含数千张各种各样的目标的图像,以及每张图像中内容的标签,还需要一堆不是汽车的图像,因为我们需要区分这两个类别。
本教程将创建一个模型,来检测图像中是否包含汽车,也可以用于检测其他目标物体!只需将下面的源 URL 更改为包含另一种目标图像的源 URL。
下载 斯坦福汽车数据集,该数据集包含 16,185 个彩色汽车图像。还需要不是汽车的随机物体的图像,这里我们使用的是 COCO 2017 验证集(比完整的训练集小,因此下载速度快,在完整数据集上进行训练效果更佳)。注意,COCO 2017 数据集中有一些汽车图像,但是数量很少,不会对结果产生 影响 - 只会稍微降低准确度。
使用 TensorFlow 数据加载器程序,并改为手动操作,以确保能够轻松更改正在使用的数据集。最终得到以下文件层次结构:
/root
├── images
│ ├── object
│ │ ├── 000001.jpg
│ │ │ ...
│ │ └── 016185.jpg
│ ├── object.tgz
│ ├── random
│ │ ├── 000000000139.jpg
│ │ │ ...
│ │ └── 000000581781.jpg
│ └── random.zip
还应该注意到,斯坦福汽车有 8k 图像,而 COCO 2017 验证集是 5k 图像——并不是对半分割!若愿意可在训练期间对这些类进行不同的加权进行纠正,不纠正仍然可以进行有效训练。下载斯坦福汽车数据集大约需要 2分钟,而 COCO 2017 验证集需要 1分钟。
import os
import shutil
import urllib.request
# 下载数据集
os.makedirs(f"{FOLDER}/downloads")
os.makedirs(f"{FOLDER}/images")
urllib.request.urlretrieve(
"https://data.deepai.org/stanfordcars.zip", f"{FOLDER}/downloads/target.zip"
)
urllib.request.urlretrieve(
"http://images.cocodataset.org/zips/val2017.zip", f"{FOLDER}/downloads/random.zip"
)
# 解压并重命名它们的文件夹
shutil.unpack_archive(f"{FOLDER}/downloads/target.zip", f"{FOLDER}/downloads")
shutil.unpack_archive(f"{FOLDER}/downloads/random.zip", f"{FOLDER}/downloads")
shutil.move(f"{FOLDER}/downloads/cars_train/cars_train", f"{FOLDER}/images/target")
shutil.move(f"{FOLDER}/downloads/val2017", f"{FOLDER}/images/random")
输出结果:
'/tmp/tmpijas024t/images/random'
加载数据
目前,数据以各种大小的 JPG 文件形式存储在磁盘上。要使用其进行训练,必须将图像加载到内存中,并调整为 64x64,然后转换为原始的、未压缩的数据。可用 Keras 的 image_dataset_from_directory
来解决该问题,它加载图像时,每个像素值都是从 0 到 255 的浮点数。
从子目录结构中得知 /objects
中的图像是一类,而 /random
中的图像是另一类。设置 label_mode='categorical'
让 Keras 将这些转换为分类标签(一个 2x1 向量),对目标类的对象来说是 [1, 0]
,对于其他任何东西来说是 [0, 1]
向量,此外,还将设置 shuffle=True
以随机化示例的顺序。
将样本分组以加快训练速度,设置 batch_size = 32
。
最后,在机器学习中,通常希望输入是小数字。因此使用 Rescaling
层来更改图像,使每个像素都是 0.0
到 1.0
之间的浮点数,而不是 0
到 255
。注意,因为要使用 lambda
函数 ,所以不要重新调整分类标签。
IMAGE_SIZE = (64, 64, 3)
unscaled_dataset = tf.keras.utils.image_dataset_from_directory(
f"{FOLDER}/images",
batch_size=32,
shuffle=True,
label_mode="categorical",
image_size=IMAGE_SIZE[0:2],
)
rescale = tf.keras.layers.Rescaling(scale=1.0 / 255)
full_dataset = unscaled_dataset.map(lambda im, lbl: (rescale(im), lbl))
输出结果:
Found 13144 files belonging to 2 classes.
数据集中有什么?
在将数据集提供给神经网络之前,应对其进行快速验证。数据是否能正确转换?标签是否合适?目标物体与其他物体的比例是多少?可以用 matplotlib
从数据集中显示一些示例:
import matplotlib.pyplot as plt
num_target_class = len(os.listdir(f"{FOLDER}/images/target/"))
num_random_class = len(os.listdir(f"{FOLDER}/images/random/"))
print(f"{FOLDER}/images/target contains {num_target_class} images")
print(f"{FOLDER}/images/random contains {num_random_class} images")
# 显示一些样本及其标签
SAMPLES_TO_SHOW = 10
plt.figure(figsize=(20, 10))
for i, (image, label) in enumerate(unscaled_dataset.unbatch()):
if i >= SAMPLES_TO_SHOW:
break
ax = plt.subplot(1, SAMPLES_TO_SHOW, i + 1)
plt.imshow(image.numpy().astype("uint8"))
plt.title(list(label.numpy()))
plt.axis("off")
输出结果:
/tmp/tmpijas024t/images/target contains 8144 images
/tmp/tmpijas024t/images/random contains 5000 images
验证准确度
开发模型时经常要检查它的准确度(例如,看看它在训练期间是否有所改进)。如何做到这一点?可以在所有数据上训练模型,然后让它对相同的数据进行分类。但是,模型可以通过记住所有样本来作弊,这会使其看起来具有非常高的准确性,但实际上表现非常糟糕。在实践中,这种“记忆”被称为过拟合。
为防止这种情况,将留出一些数据(20%)作为验证集,用来检查模型的准确性。
num_batches = len(full_dataset)
train_dataset = full_dataset.take(int(num_batches * 0.8))
validation_dataset = full_dataset.skip(len(train_dataset))