Qualcomm AI Hub-示例（一）编译模型

文章介绍

Qualcomm® AI Hub支持Pytorch、ONNX、AIMET、Tensorflow等主流AI框架训练的模型，本文介绍了如何使用AI Hub相关的SDK接口把这些模型编译成TensorFlow Lite和Qualcomm AI Engine Direct (QNN)。原文可以参阅 Compiling Models

编译模型

Qualcomm AI Hub支持编译使用以下AI框架训练的模型：

• PyTorch
• ONNX
• AI Model Efficiency Toolkit (AIMET) quantized models.
• TensorFlow (via ONNX)

可以为以下目标运行时编译上述任何模型：

• TensorFlow Lite
• Qualcomm AI Engine Direct (QNN)

将PyTorch编译为TensorFlow Lite

要编译PyTorch模型，我们必须首先使用PyTorch中的jit.trace方法在内存中生成TorchScript模型。一旦触发，就可以使用submit_compile_job() API编译模型。

from typing import Tuple

import torch

import torchvision

import qai_hub as hub

# Using pre-trained MobileNet

torch_model = torchvision.models.mobilenet_v2(pretrained=True)

torch_model.eval()

# Trace model

input_shape: Tuple[int, ...] = (1, 3, 224, 224)

example_input = torch.rand(input_shape)

pt_model = torch.jit.trace(torch_model, example_input)

# Compile model on a specific device

compile_job = hub.submit_compile_job(

    pt_model,

    name="MyMobileNet",

    device=hub.Device("Samsung Galaxy S23 Ultra"),

    input_specs=dict(image=input_shape),

)

assert isinstance(compile_job, hub.CompileJob)

编译TorchScript模型

如果您已经保存了跟踪或脚本化的torch模型（使用torch.jit.save保存），则可以直接提交。我们将以mobilenet_v2.pt为例。

import qai_hub as hub

# Compile a model

compile_job = hub.submit_compile_job(

    model="mobilenet_v2.pt",

    device=hub.Device("Samsung Galaxy S23 Ultra"),

    input_specs=dict(image=(1, 3, 224, 224)),

)

assert isinstance(compile_job, hub.CompileJob)

# Profile the compiled model

profile_job = hub.submit_profile_job(

    model=compile_job.get_target_model(),

    device=hub.Device("Samsung Galaxy S23 Ultra"),

)

assert isinstance(profile_job, hub.ProfileJob)

将TorchScript编译为QNN模型库

Qualcomm AI Hub支持对TorchScript到QNN模型库进行编译和分析。在本例中，我们将使用mobilenet_v2.pt并将其编译为ARM64 Android平台（aarch64_Android）的QNN模型库（.so文件）。

import qai_hub as hub

# Compile a model to a QNN Model Library

compile_job = hub.submit_compile_job(

    model="mobilenet_v2.pt",

    device=hub.Device("Samsung Galaxy S23"),

    options="--target_runtime qnn_lib_aarch64_android",

    input_specs=dict(image=(1, 3, 224, 224)),

)

assert isinstance(compile_job, hub.CompileJob)

返回值是CompileJob的一个实例。看看这里的例子，了解如何为Snapdragon® Neural Processing Unit（NPU）编译此模型。

将ONNX模型编译为TensorFlow Lite或QNN

Qualcomm AI Hub还支持将ONNX模型编译为TensorFlow Lite或QNN模型库。我们将以mobilenet_v2.onnx为例。

import qai_hub as hub

# Compile a model to TensorFlow Lite

compile_job = hub.submit_compile_job(

    model="mobilenet_v2.onnx",

    device=hub.Device("Samsung Galaxy S23 Ultra"),

)

assert isinstance(compile_job, hub.CompileJob)

# Compile a model to a QNN Model Library

compile_job = hub.submit_compile_job(

    model="mobilenet_v2.onnx",

    device=hub.Device("Samsung Galaxy S23"),

    options="--target_runtime qnn_lib_aarch64_android",

)

assert isinstance(compile_job, hub.CompileJob)

将AIMET量化的模型编译为TensorFlow Lite或QNN

AI Model Efficiency Toolkit（AIMET）是一个开源库，为训练神经网络模型提供先进的模型量化和压缩技术。AIMET的QuantizationSimModel可以导出到以下其中一个：

• ONNX模型（.onx）和带有量化参数的编码文件（.encodings）。(推荐)
• TorchScript（.pt）和具有量化参数的编码文件（.encodings）。

要使用这些模型，请创建一个名称为.amet的目录。它应该包含一个.pt或.onx模型和相应的编码文件。

<model>.aimet

   |-- <model>.onnx  or .pt file

   +-- <model>.encodings

其中＜model＞可以是任何名称。

让我们以mobilenet_v2_onnx.aimet.zip为例。解压到mobilenet_v2_onnx.amet目录后，我们可以通过

import qai_hub as hub

# Compile to TensorFlow Lite

compile_job = hub.submit_compile_job(

    model="mobilenet_v2_onnx.aimet",

    device=hub.Device("Samsung Galaxy S23"),

)

assert isinstance(compile_job, hub.CompileJob)

# Compile to a QNN Model Library

compile_job = hub.submit_compile_job(

    model="mobilenet_v2_onnx.aimet",

    device=hub.Device("Samsung Galaxy S23"),

    options="--target_runtime qnn_lib_aarch64_android --quantize_full_type int8",

)

assert isinstance(compile_job, hub.CompileJob)

作者：高通工程师，戴忠忠（Zhongzhong Dai）