5. 在 Arduino 上为 microTVM 训练视觉模型
单击 此处 下载完整的示例代码
作者:Gavin Uberti
本教程介绍如何训练 MobileNetV1 模型以适应嵌入式设备,以及如何使用 TVM 将这些模型部署到 Arduino。
背景简介
构建物联网设备时,通常想让它们能够看到并理解它们周围的世界。可以采取多种形式,但通常设备也会想知道某种物体是否在其视野中。
例如,安全摄像头可能会寻找人,因此它可以决定是否将视频保存到内存中。红绿灯可能会寻找汽车,这样它就可以判断哪个信号灯应该首先改变。或者森林相机可能会寻找一种动物,从而估计动物种群的数量。
为使这些设备价格合理,我们希望为这些设备配置一个低成本处理器,如 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))
加载数据
在过 去的十年中,卷积神经网络 已被广泛用于图像分类任务。EfficientNet V2 等最先进的模型的图像分类效果甚至优于人类。然而,这些模型有数千万个参数,它们无法运行在便宜的监控摄像头计算机上。
应用程序的准确度达到 90% 就足够了。因此可以使用更旧版本更小的 MobileNet V1 架构。但这仍然不够小,默认情况下,对于具有 224x224 输入和 alpha 1.0 的 MobileNet V1,仅存储就需要大约 50 MB。为了减小模型的大小,可调节三个 knob。
首先,可以将输入图像的大小从 224x224 减小到 96x96 或 64x64,Keras 可以轻松做到这一点。还可以将模型的 alpha 从 1.0 降低到 0.25,这使得网络的宽度(和过滤器的数量)缩小了四倍。如果真的空间有限,可以通过让模型拍摄灰度图像而不是 RGB 图像来减少通道数量。
在本教程中,使用 RGB 64x64 输入图像和 alpha 0.25。效果虽然不是很理想,但它允许完成的模型适合 192 KB 的 RAM,同时仍然允许使用官方 TensorFlow 源模型执行迁移学习(如果使用 alpha <0.25 或灰度输入,无法做到这样)。
什么是迁移学习?
深度学习长期以来一直 主导图像分类,但训练神经网络需要大量时间。当一个神经网络“从头开始”训练时,起初参数会随机初始 化,缓慢地学习如何区分图像。
从一个已经擅长特定任务的神经网络开始迁移学习,在此示例中,该任务是对 ImageNet 数据集 中的图像进行分类。这意味着神经网络已经具有一些目标检测。
这对于像 MobileNet 这样的图像处理神经网络特别有效,在实践中,模型的卷积层(即前 90% 的层)用于识别 line 和 shape 等低级特征——只有最后几个全连接层用于确定,这些 shape 如何构成网络要检测的目标。
可以通过使用在 ImageNet 上训练的 MobileNet 模型开始训练,该模型已经知道如何识别线条和形状。从这个预训练模型中删除最后几层,并添加自己的最终层。在汽车与非汽车数据集上训练这个联合模型,以调优第一层并从头开始训练最后一层。这种训练已经部分训练过的模型称为微调。
源 MobileNets 模型(用于迁移学习) 已被 TensorFlow 人员预训练,因此可以下载最接近于我们想要的版本( 128x128 输入模型,深度为 0.25)。
os.makedirs(f"{FOLDER}/models")
WEIGHTS_PATH = f"{FOLDER}/models/mobilenet_2_5_128_tf.h5"
urllib.request.urlretrieve(
"https://storage.googleapis.com/tensorflow/keras-applications/mobilenet/mobilenet_2_5_128_tf.h5",
WEIGHTS_PATH,
)
pretrained = tf.keras.applications.MobileNet(
input_shape=IMAGE_SIZE, weights=WEIGHTS_PATH, alpha=0.25
)
修改网络
如上所述,预训练模型旨在对 1,000 个 ImageNet 类别进行分类,但我们希望将其转换为对汽车进行分类。由于只有最后几层是特定于任务的,因此切断原始模型的最后五层,通过执行 respape、dropout、flatten 和 softmax 操作来构建自己的模型的最后几层。
model = tf.keras.models.Sequential()
model.add(tf.keras.layers.InputLayer(input_shape=IMAGE_SIZE))
model.add(tf.keras.Model(inputs=pretrained.inputs, outputs=pretrained.layers[-5].output))
model.add(tf.keras.layers.Reshape((-1,)))
model.add(tf.keras.layers.Dropout(0.1))
model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(2, activation="softmax"))
微调网络
在训练神经网络时,必须设置学习率来控制网络学习速度。太慢和太快都会导致学习效果不好。通常对于 Adam(正在使用的优化器)来说,0.001
是一个相当不错的学习率(并且是 原始论文 中推荐的)。在本示例中,0.0005
效果更好。
将之前的验证集传递给 model.fit
,在每次训练时评估模型性能,并能够跟踪模型性能是如何提升的。训练完成后,模型的验证准确率应该在 0.98
左右(这意味着在验证集上训练 100 次,其中 98 次都是预测正确的)。
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=0.0005),
loss="categorical_crossentropy",
metrics=["accuracy"],
)
model.fit(train_dataset, validation_data=validation_dataset, epochs=3, verbose=2)
输出结果:
Epoch 1/3
328/328 - 56s - loss: 0.2151 - accuracy: 0.9280 - val_loss: 0.1213 - val_accuracy: 0.9615
Epoch 2/3
328/328 - 53s - loss: 0.1029 - accuracy: 0.9623 - val_loss: 0.1111 - val_accuracy: 0.9626
Epoch 3/3
328/328 - 52s - loss: 0.0685 - accuracy: 0.9750 - val_loss: 0.1541 - val_accuracy: 0.9505
<keras.callbacks.History object at 0x7efef1110b90>
量化
通过改变输入维度, 以及移除底层,将模型减少到只有 219k 个参数,每个参数都是 float32
类型,占用 4 个字节,因此模型将占用将近 1 MB!
此外,硬件可能没有内置对浮点数的支持,虽然大多数高内存 Arduino(如 Nano 33 BLE)确实有硬件支持,但其他一些(如 Arduino Due)则没有。在任何没有专用硬件支持的板上,浮点乘法都会非常慢。
为解决这两个问题可以将模型量化,把权重表示为八位整数,为获得最佳性能,TensorFlow 会跟踪模型中每个神经元的激活方式,因此可以得出,如何最准确地用整数运算,来模拟神经元的原始激活。
可以创建一个具有代表性的数据集来帮助 TensorFlow 实现——原始数据集的一个子集,用于跟踪这些神经元的激活方式。然后将其传递给带有 Optimize
标志的 TFLiteConverter
(Keras 本身不支持量化),以告诉 TFLite 执行转换。默认情况下,TFLite 将模型的输入和输出保持为浮点数,因此必须明确告诉它避免这种行为。
def representative_dataset():
for image_batch, label_batch in full_dataset.take(10):
yield [image_batch]
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
quantized_model = converter.convert()