使用NVIDIA TAO工具包优化Arm Ethos-U NPUs的AI模型

本文翻译自：《Optimizing AI models for Arm Ethos-U NPUs using the NVIDIA TAO Toolkit》

作者：Amogh Dabholkar

优化实现推理吞吐量增加4倍，内存减少3倍

边缘端的AI大规模应用带来了多个优势，包括降低延迟、增强隐私和成本效益。Arm一直处于这一领域的前沿，专注于在其Cortex-A和Cortex-M CPU以及Ethos-U NPU上提供先进的边缘AI功能。然而，这个领域仍在迅速扩张，对于希望在数十亿边缘设备上轻松部署的开发者来说，这带来了挑战。

其中一个挑战是为边缘设备开发深度学习模型，因为开发者需要处理有限的资源，如存储、内存和计算能力，同时平衡模型准确性和延迟或帧率等运行时指标。一个为更强大平台设计的现成模型在部署到资源受限的平台上可能会很慢，甚至无法运行。

TAO工具包是由NVIDIA在Tensorflow和PyTorch基础上开发的低代码开源工具，用于抽象训练深度学习模型的复杂性。它拥有广泛的计算机视觉应用的预训练模型库，以便进行迁移学习，并提供通道剪枝和量化感知训练等形式的即插即用模型优化，从而帮助开发更轻量的模型。

图片

在本博客中，我们：

• 使用TAO获取在ImageNet数据上预训练的MobilenetV2模型，并在Visual Wake Words数据集上进行微调。
• 利用TAO中的通道剪枝来优化模型，减小模型尺寸并提高推理吞吐量。
• 比较这些模型在Arm Ethos-U NPU上运行的性能。
• 演示在最小牺牲模型准确性的情况下在Arm Ethos-U NPU上实现更快的模型运行时间。

如果您想了解有关使用其他类型的模型优化技术（如随机剪枝和聚类）在Arm Ethos-U NPU上的优势，请阅读本博客。

前提条件：

我们假设您：

• 熟悉如何训练深度学习模型。
• 熟悉Arm Vela编译器和ML Inference Advisor。
• 曾在Ethos-U NPU上运行神经网络。ml-embedded-evaluation-kit允许您轻松在Ethos-U NPU上执行神经网络。

完整的可执行代码将在GitHub上提供。

运行NVIDIA TAO设置：

所需的设置非常简单，包括以下步骤：

1. 安装Docker并按照https://docs.docker.com/engine/install/linux-postinstall/进行设置。
2. 安装NVIDIA容器工具包。
3. 设置NGC帐户并获取NGC API密钥。
4. 在终端上，使用docker login nvcr.io登录到NGC Docker Registry。这将使TAO能够根据需要拉取必要的Docker映像。

图片

1. 设置一个conda环境，以进一步隔离Python包的依赖关系。
2. 从NGC Registry下载最新的TAO软件包，然后您可以开始使用Jupyter Notebook

wget --content-disposition https://api.ngc.nvidia.com/v2/resources/nvidia/tao/tao-getting-started/versions/5.0.0/zip -O getting_started_v5.0.0.zipunzip -u getting_started_v5.0.0.zip  -d ./getting_started_v5.0.0 && rm -rf getting_started_v5.0.0.zip && cd ./getting_started_v5.0.0

如需更详细的说明，请参考官方NVIDIA设置说明页面，点击此处：

https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_quick_start_guide.html

设置完成后，您可以从Arm ML-Examples存储库中下载我们的Jupyter笔记本，并将其保存到最近下载的TAO文件夹内的以下路径中：tao-getting-started_v5.0.0/notebooks/tao_launcher_starter_kit/classification_tf2/tao_voc/ 

获取TAO Toolkit的基础模型

TAO Toolkit具有广泛的模型可供使用，作为其模型仓库的一部分，这些模型可以轻松下载并用于各种应用。您可以使用以下方式使用NGC CLI来获取可用的预训练模型列表。

!ngc registry model list nvidia/tao/pretrained_classification:*

我们将下载在ImageNet上预训练的MobilenetV2模型，并将其用于训练我们自己的下游任务。

# Pull pretrained model from NGC!ngc registry model download-version nvidia/tao/pretrained_classification:mobilenet_v2 --dest $LOCAL_EXPERIMENT_DIR/pretrained_mobilenet_v2

进行自定义任务训练

一旦您下载了预训练模型，只要数据集满足以下格式，您就可以对其进行微调：

图片

每个类别名称文件夹应包含与该类别对应的图像。相同的类别名称文件夹应存在于images_test、images_train和images_val中。txt文件包含所有类别的名称（每个名称单独一行）。

根据这些准则，我们可以将MobileNetV2模型迁移到Visual Wake Words数据集上。Visual Wake Words数据集是从COCO数据集派生而来，用于训练模型以检测图像帧中是否存在人物，这对物联网设备尤其重要。这是一个包含两个类别的图像分类问题：

图片

我们使用TAO中的以下命令行来训练模型：

print("To run this training in data parallelism using multiple GPU's, please uncomment the line below and "    " update the --gpus parameter to the number of GPU's you wish to use.")!tao model classification_tf2 train -e $SPECS_DIR/spec.yaml --gpus 2

训练完成后，基础模型的评估准确率为90.33％。

在Arm Ethos-U NPU上运行基础模型的性能

要将模型部署到Arm Ethos-U NPU上，我们需要对模型进行INT8量化。NVIDIA提供的所有模型都是使用EFF进行编码的。NVIDIA Exchange File Format (EFF) 的创建旨在促进不同NVIDIA深度学习框架和工具之间的交换和互操作性。我们将使用下面所示的decode_eff()函数首先将模型解码回TensorFlow格式，然后我们将使用以下代码进行训练后量化（PTQ），并获得一个INT8 tflite模型。

def representative_dataset_gen():    root_path = '/home/amodab01/v5.0/mobilenetV2/visualwake/data/train'    categories =  os.listdir(root_path)    x = []    img_path = os.path.join(root_path, categories[0])    images = os.listdir(img_path)    for j in range(100):        img = cv2.imread(os.path.join(img_path, images[j]))        img = cv2.resize(img, (224,224))        img = img/255.0        img = img.astype(np.float32)        x.append(img)    x = np.array(x)    train_data = tf.data.Dataset.from_tensor_slices(x)    for i in train_data.batch(1).take(5):        yield [i]

def decode_eff(eff_model_path, enc_key=None):    """Decode EFF to saved_model directory.    Args:        eff_model_path (str): Path to eff model        enc_key (str, optional): Encryption key. Defaults to None.    Returns:        str: Path to the saved_model    """    # Decrypt EFF    eff_filename = os.path.basename(eff_model_path)    eff_art = Archive.restore_artifact(        restore_path=eff_model_path,        artifact_name=eff_filename,        passphrase=enc_key)    zip_path = eff_art.get_handle()    # Unzip    saved_model_path = os.path.dirname(zip_path)    with zipfile.ZipFile(zip_path, "r") as zip_file:        zip_file.extractall(saved_model_path)    return saved_model_path

input_model_file = '/home/amodab01/v5.0/mobilenetV2/visualwake/classification_tf2/output/train/mobilenet_v2_bn_070.tlt'output_model_file = '/home/amodab01/v5.0/mobilenetV2/visualwake/classification_tf2/output/int81/model.tflite'key = 'tlt'if os.path.isdir(input_model_file):    print("Model provided is a saved model directory at {}".format(input_model_file))    saved_model = input_model_fileelse:    saved_model = decode_eff(        input_model_file,        enc_key=key    )print("Converting the saved model to tflite model.")converter = tf.lite.TFLiteConverter.from_saved_model(    saved_model,    signature_keys=["serving_default"],)converter.optimizations = [tf.lite.Optimize.DEFAULT]converter.target_spec.supported_ops = [    tf.lite.OpsSet.TFLITE_BUILTINS_INT8, # enable TensorFlow Lite ops.    tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.]converter.inference_input_type = tf.int8converter.inference_output_type = tf.int8converter.representative_dataset = representative_dataset_gentflite_model = converter.convert()model_root = os.path.dirname(output_model_file)if not os.path.exists(model_root):    os.makedirs(model_root)print("Writing out the tflite model.")with open(output_model_file, "wb") as tflite_file:    model_size = tflite_file.write(tflite_model)print(f"TFLite model of size {model_size//MB} MB was written to {output_model_file}")

我们使用Vela编译器以及Corstone-300 Fixed Virtual Platform来获取在Arm Ethos-U NPU上运行tflite模型的性能数据。

Vela是由Arm开发的，用于将tflite模型编译成优化版本，可以在包含Arm Ethos-U NPU的嵌入式系统上运行。它是一个Python包，可以使用以下方式安装：

pip install ethos-u-vela

更多详细信息可以在这里找到：

https://review.mlplatform.org/plugins/gitiles/ml/ethos-u/ethos-u-vela

Corstone-300是一个用于Cortex-M微控制器的Ethos-U NPU的周期近似模拟器。有关Corstone-300的更多信息可以在这里找到：

https://developer.arm.com/Processors/Corstone-300。

开发者可以通过ML Inference Adviser项目（https://pypi.org/project/mlia/）获得它。可以使用以下命令进行安装：

pip install mlia
mlia-backend install corstone-300

Vela基于简化的假设进行估算，这就是为什么它与在Corstone-300上测得的结果不完全相同的原因。

SRAM使用情况也基于Corstone-300的估算。我们在Corstone-300 FPV中使用以下设置：

mlia check --output-dir tao-mnv2 --performance -t ethos-u55-256 tao/MobilenetV2.tflite

这个配置对应于具有256个MAC引擎的Arm Ethos-U55 NPU。 

从TAO Toolkit获取通道剪枝模型

对于TF2，TAO Toolkit提供了以下通道剪枝选项及参数：

图片

通道剪枝旨在删除每层中不重要的通道，以便模型可以在对准确性的最小影响下缩小。要开始，我们将首先尝试使用剪枝阈值为0.5，这将删除每层约50%的通道，并使用其他参数的默认值。请注意，通道剪枝会减少每层的输入和输出通道数量，以匹配大小，因此由其他因素（如粒度和最小通道数）控制的结果模型将小于50%。

要对模型进行剪枝，我们使用以下命令：

!tao model classification_tf2 prune -e $SPECS_DIR/spec.yaml

使用0.5的阈值，剪枝后的模型大小约比原始模型小了大约4倍，但可能会略微降低准确性。这是因为一些先前有帮助的权重可能已被剔除。建议重新使用相同数据集对这个剪枝后的模型进行重新训练，以恢复准确性。重新训练后，我们获得了90.35%的评估准确性。与基准密集模型相比，重新训练已经恢复了所有丢失的准确性。

为了获得更快速的模型以适应更小的延迟预算，我们可以使用更高的剪枝阈值0.68，这会删除每层大约68%的通道。剪枝后的模型约小了大约10倍，重新训练后，模型的评估准确性为90.17%，几乎恢复了之前丢失的准确性，同时模型更小。

在下图中，我们使用Netron可视化模型，并在剪枝前后将它们的图结构并排进行比较。请注意通道的数量（每个红色框的最后一个维度）已经减少。完整模型显示在左侧，具有0.68剪枝阈值的剪枝模型显示在右侧。

剪枝模型在Ethos-U NPU上的性能

要在Ethos-U上部署模型，我们需要使用训练后量化将模型量化为INT8。与密集模型类似，我们使用前一部分提供的代码块来获取INT8 tflite模型，这些模型可以与Vela一起编译，并得到以下性能估算。

图片

AutoML

此外，TAO Toolkit提供了两个AutoML算法——Hyperband和Bayesian作为API服务的一部分，这些算法可以用于自动调整特定模型和数据集对的超参数。我们将在未来的博客文章中介绍如何使用AutoML功能，以进一步提高模型的准确性，并比较每个算法的权衡和取舍。

结论

本博客介绍了如何使用NVIDIA TAO Toolkit中提供的预训练模型，将其适应于自定义数据集和用例，然后使用TAO中的通道剪枝功能获取符合延迟要求并在Arm Ethos-U NPU上获得更好性能的模型。使用现成的预训练模型使用户能够快速针对较小的数据集进行下游任务的微调，同时仍能够实现高准确性。TAO Toolkit简化了这一过程，并提供了良好的优化选项，使用户能够在不牺牲模型准确性太多的情况下获得3倍到4倍的性能和吞吐量。它还提供了将模型部署到高性能Arm Ethos-U NPU的途径，为在Arm边缘部署机器学习模型开启了无限的机会。我们鼓励开发人员尝试NVIDIA TAO Toolkit，并将其用于优化在Arm硬件上部署的模型。