初探TVM--通过TVM的python接口编译优化resnet50

初探TVM--通过TVM的python接口编译优化resnet50

通过TVM的python接口编译优化resnet50

• 通过TVM的python接口编译优化resnet50
• • 下载并加在onnx模型
  • 下载和预处理图片
  • 用relay接口编译模型
  • 基于TVM运行时组件执行编译后的模型
  • 基础优化的数据

通过TVM的python接口编译优化resnet50

在上一章的教程中^{1，我们通过tvmc这个command line工具优化并且调优（tune）了一个预训练好的视觉模型–resnet50 v2. 不过tvm本身也有一套基于python的API，他们在优化深度学习模型的工作上提供了强大的灵活性。
在本章的教程中，我们会继续使用在tvmc教程中的背景，但是主要使用python API来完成任务，而不是使用TVMC工具。 在本章节的教程中，我们会使用TVM的python API实现以下任务功能：}

• 基于tvm运行时组件，编译一个预训练好的resnet 50 v2模型
• 通过编译过的模型，在一张真实图片上跑出运算结果
• 在CPU上对模型调优
• 编译基于调优过的模型
• 在真实图片上，运行调优过的模型，获取输出

本章节的目标是对tvm的使用场景有一个大体的理解，并且能够知道怎样使用tvm的python API来完成一个模型编译和优化的任务。

使用python API写脚本的话，首先会需要import很多必须的库，例如onnx、numpy等等。

import onnx
from tvm.contrib.download import download_testdata
from PIL import Image
import numpy as np
import tvm.relay as relay
import tvm
from tvm.contrib import graph_executor

下载并加在onnx模型

在本章节中，我们仍然使用resnet 50 v2，这是一个有50层卷积的图片分类模型。模型通过超过100万张图片和1000中不同分类做过预训练，需要输入的图片分辨率为224x224。由于之前内容对resnet50做过介绍，这里就不再多说了。
TVM封装了一个库可以用来下载预训练的模型，只需要提供模型地址，类型等信息，TVM可以有API来完成模型下载和保存。

model_url = "".join(
    [
        "https://github.com/onnx/models/raw/",
        "master/vision/classification/resnet/model/",
        "resnet50-v2-7.onnx",
    ]
)

model_path = download_testdata(model_url, "resnet50-v2-7.onnx", module="onnx")
onnx_model = onnx.load(model_path)

下载和预处理图片

和前面章节一样，我们还会再把这只可爱的小猫咪拉出来分类（写到这里，又想我们家提米了）。

下载和预处理的代码在下面啦：

img_url = "https://s3.amazonaws.com/model-server/inputs/kitten.jpg"
img_path = download_testdata(img_url, "imagenet_cat.png", module="data")

# Resize it to 224x224
resized_image = Image.open(img_path).resize((224, 224))
img_data = np.asarray(resized_image).astype("float32")

# Our input image is in HWC layout while onNX expects CHW input, so convert the array
img_data = np.transpose(img_data, (2, 0, 1))

# Normalize according to the ImageNet input specification
imagenet_mean = np.array([0.485, 0.456, 0.406]).reshape((3, 1, 1))
imagenet_stddev = np.array([0.229, 0.224, 0.225]).reshape((3, 1, 1))
norm_img_data = (img_data / 255 - imagenet_mean) / imagenet_stddev

# Add the batch dimension, as we are expecting 4-dimensional input: NCHW.
img_data = np.expand_dims(norm_img_data, axis=0)

用relay接口编译模型

接下来我们就可以编译这个resnet50的模型了。首先使用relay的onnx模块导入，然后用标准优化流程编译，最后会创建出一个TVM的图运行时模块：

target = "llvm"
# The input name may vary across model types. You can use a tool
# like Netron to check input names
input_name = "data"
shape_dict = {input_name: img_data.shape}

mod, params = relay.frontend.from_onnx(onnx_model, shape_dict)

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

dev = tvm.device(str(target), 0)
module = graph_executor.GraphModule(lib["default"](dev))

如果能够在指定target时给出准确的平台信息，就能够获得更好的性能，因为TVM内部会利用平台特性，做出相应的优化策略，例如target = "llvm -mcpu=skylake" 或者 target="llvm -mcpu=skylake-avx512"，这样就可以利用X86的avx512指令集优化。

基于TVM运行时组件执行编译后的模型

在编译好模型后，就可以用tvm 运行时组件做推理运算了。

dtype = "float32"
module.set_input(input_name, img_data)
module.run()
output_shape = (1, 1000)
tvm_output = module.get_output(0, tvm.nd.empty(output_shape)).numpy()

不知道这里换成float16会不会也可以

基础优化的数据

在调优模型之前，我们先看一下基础的模型的优化数据。为了是测试准确，我们多次运行模型，计算平均计算时间。

import timeit

timing_number = 10
timing_repeat = 10
unoptimized = (
    np.array(timeit.Timer(lambda: module.run()).repeat(repeat=timing_repeat, number=timing_number))
    * 1000
    / timing_number
)
unoptimized = {
    "mean": np.mean(unoptimized),
    "median": np.median(unoptimized),
    "std": np.std(unoptimized),
}

print(unoptimized)

---

1. 使用TVMC优化resnet50【1】【2】 ↩︎