Ctrl+K

使用 TVM 部署框架: 预量化模型-第3部分(TFLite)

%%shell
# Installs the latest dev build of TVM from PyPI. If you wish to build
# from source, see https://tvm.apache.org/docs/install/from_source.html
pip install apache-tvm --pre

使用 TVM 部署框架: 预量化模型-第3部分(TFLite)#

Author: Siju Samuel

欢迎来到部署框架的第3部分——使用 TVM 预量化模型教程。

在这一部分中，将从量化的 TFLite graph 开始，然后通过 TVM 编译和执行它。

有关使用 TFLite 量化模型的更多细节，建议读者阅读 转换量化模型。

TFLite 模型可以从这个 hosted_models 下载。

开始之前，需要先安装 Tensorflow 和 TFLite 包。

# install tensorflow and tflite
pip install tensorflow==2.1.0
pip install tflite==2.1.0

现在请检查 TFLite 包是否安装成功，python -c "import tflite"

必需的导入#

import os

import numpy as np
import tflite

import tvm
from tvm import relay

下载预训练的量化 TFLite 模型#

# Download mobilenet V2 TFLite model provided by Google
from tvm.contrib.download import download_testdata

model_url = (
    "https://storage.googleapis.com/download.tensorflow.org/models/"
    "tflite_11_05_08/mobilenet_v2_1.0_224_quant.tgz"
)

# Download model tar file and extract it to get mobilenet_v2_1.0_224.tflite
model_path = download_testdata(
    model_url, "mobilenet_v2_1.0_224_quant.tgz", module=["tf", "official"]
)
model_dir = os.path.dirname(model_path)

Utils 用于下载和解压zip文件#

def extract(path):
    import tarfile

    if path.endswith("tgz") or path.endswith("gz"):
        dir_path = os.path.dirname(path)
        tar = tarfile.open(path)
        tar.extractall(path=dir_path)
        tar.close()
    else:
        raise RuntimeError("Could not decompress the file: " + path)


extract(model_path)

加载测试图片#

获取真实图像进行端到端（e2e）测试#

def get_real_image(im_height, im_width):
    from PIL import Image

    repo_base = "https://github.com/dmlc/web-data/raw/main/tensorflow/models/InceptionV1/"
    img_name = "elephant-299.jpg"
    image_url = os.path.join(repo_base, img_name)
    img_path = download_testdata(image_url, img_name, module="data")
    image = Image.open(img_path).resize((im_height, im_width))
    x = np.array(image).astype("uint8")
    data = np.reshape(x, (1, im_height, im_width, 3))
    return data


data = get_real_image(224, 224)

加载 tflite 模型#

现在我们可以打开 mobilenet_v2_1.0_224.tflite

tflite_model_file = os.path.join(model_dir, "mobilenet_v2_1.0_224_quant.tflite")
tflite_model_buf = open(tflite_model_file, "rb").read()

# Get TFLite model from buffer
try:
    import tflite

    tflite_model = tflite.Model.GetRootAsModel(tflite_model_buf, 0)
except AttributeError:
    import tflite.Model

    tflite_model = tflite.Model.Model.GetRootAsModel(tflite_model_buf, 0)

让我们运行 TFLite 预量化模型推断并获得 TFLite 预测。

def run_tflite_model(tflite_model_buf, input_data):
    """Generic function to execute TFLite"""
    try:
        from tensorflow import lite as interpreter_wrapper
    except ImportError:
        from tensorflow.contrib import lite as interpreter_wrapper

    input_data = input_data if isinstance(input_data, list) else [input_data]

    interpreter = interpreter_wrapper.Interpreter(model_content=tflite_model_buf)
    interpreter.allocate_tensors()

    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    # set input
    assert len(input_data) == len(input_details)
    for i in range(len(input_details)):
        interpreter.set_tensor(input_details[i]["index"], input_data[i])

    # Run
    interpreter.invoke()

    # get output
    tflite_output = list()
    for i in range(len(output_details)):
        tflite_output.append(interpreter.get_tensor(output_details[i]["index"]))

    return tflite_output

让我们运行 TVM 编译的预量化模型推断并获得 TVM 预测。

def run_tvm(lib):
    from tvm.contrib import graph_executor

    rt_mod = graph_executor.GraphModule(lib["default"](tvm.cpu(0)))
    rt_mod.set_input("input", data)
    rt_mod.run()
    tvm_res = rt_mod.get_output(0).numpy()
    tvm_pred = np.squeeze(tvm_res).argsort()[-5:][::-1]
    return tvm_pred, rt_mod

TFLite 推理#

在量化模型上运行 TFLite 推理。

tflite_res = run_tflite_model(tflite_model_buf, data)
tflite_pred = np.squeeze(tflite_res).argsort()[-5:][::-1]

2023-06-08 16:07:05.224447: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.
2023-06-08 16:07:05.275854: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-06-08 16:07:06.156446: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
INFO: Created TensorFlow Lite XNNPACK delegate for CPU.

TVM 编译和推断#

我们使用 TFLite-Relay 解析器将 TFLite 预量化图转换为 Relay IR。请注意，预量化模型的前端解析器调用与 FP32 模型的前端解析器调用完全相同。我们建议你删除 print(mod) 中的注释，并检查 Relay 模块。您将看到许多 QNN 算子，如 Requantize、Quantize 和 QNN Conv2D。

dtype_dict = {"input": data.dtype.name}
shape_dict = {"input": data.shape}

mod, params = relay.frontend.from_tflite(tflite_model, shape_dict=shape_dict, dtype_dict=dtype_dict)
# print(mod)

现在让我们编译 Relay 模块。我们在这里使用“llvm”目标。请替换为您感兴趣的目标平台。

target = "llvm"
with tvm.transform.PassContext(opt_level=3):
    lib = relay.build_module.build(mod, target=target, params=params)

最后，让我们在 TVM 编译模块上调用推断。

tvm_pred, rt_mod = run_tvm(lib)

Accuracy 对比#

打印 MXNet 和 TVM 推理的 top-5 标签。检查标签，因为 TFLite 和 Relay 的重量化实现不同。这导致最终输出的数字不匹配。因此，通过标签来测试准确性。

print("TVM Top-5 labels:", tvm_pred)
print("TFLite Top-5 labels:", tflite_pred)

TVM Top-5 labels: [387 102 386 349 341]
TFLite Top-5 labels: [387 102 386 341 880]

性能度量#

文中给出了如何测量 TVM 编译模型性能的例子。

n_repeat = 100  # should be bigger to make the measurement more accurate
dev = tvm.cpu(0)
print(rt_mod.benchmark(dev, number=1, repeat=n_repeat))

Execution time summary:
 mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  
  18.9581      18.4153      33.4940      18.2674       1.7527

备注

除非硬件对快速 8 位指令有特殊支持，否则量化模型不会比 FP32 模型更快。如果没有快速的 8 位指令，TVM 在 16 位中进行量化卷积，即使模型本身是 8 位。

对于 x86，在指令集为 AVX512 的 CPU 上可以达到最好的性能。在这种情况下，TVM 为给定目标利用最快的 8 位指令。这包括对 VNNI 8 位点积指令（CascadeLake 或更新的）的支持。对于 EC2 C5.12x 大型实例，本教程的TVM延迟约为 2 ms。

在许多 TFLite 网络中，Intel conv2d NCHWc 调度比 ARM NCHW conv2d 空间包调度具有更好的端到端延迟。ARM winograd 的性能更高，但它占用的内存也更多。

此外，以下关于 CPU 性能的一般提示同样适用：

  • 将环境变量 TVM_NUM_THREADS 设置为物理核数

  • 为你的硬件选择最佳的目标，例如 “llvm -mcpu=cascadelake” 或 “llvm -mcpu=skylake-avx512” （将来会有更多带有 AVX512 的 CPU）

  • 执行自动调优

  • 为了在 ARM CPU 上获得最佳的推理性能，请根据您的设备更改目标参数并遵循 Auto-tuning a Convolutional Network for ARM CPU

目录