利用TVM读取并预测ONNX模型

You 2023-12-30 23:58:01  50895

分类专栏：资讯

利用TVM读取并预测ONNX模型

在我们成功编译并且可以在Python端正常引用TVM后，我们首先导入我们的onnx格式的模型。这里我们准备了一张飞机的图像：

这个图像在ImageNet分类中属于404: 'airliner'，也就是航空客机。

下面我们将利用TVM部署onnx模型并对这张图像进行预测。

import onnx
import time
import tvm
import numpy as np
import tvm.relay as relay
from PIL import Image

onnx_model = onnx.load('mobilenetv2.onnx')  # 导入模型

mean = [123., 117., 104.]                   # 在ImageNet上训练数据集的mean和std
std = [58.395, 57.12, 57.375]


def transform_image(image):                # 定义转化函数，将PIL格式的图像转化为格式维度的numpy格式数组
    image = image - np.array(mean)
    image /= np.array(std)
    image = np.array(image).transpose((2, 0, 1))
    image = image[np.newaxis, :].astype('float32')
    return image

img = Image.open('../datasets/images/plane.jpg').resize((224, 224)) # 这里我们将图像resize为特定大小
x = transform_image(img)

这样我们得到的x为[1,3,224,224]维度的ndarray。这个符合NCHW格式标准，也是我们通用的张量格式。

接下来我们设置目标端口llvm，也就是部署到CPU端，而这里我们使用的是TVM中的Relay IR，这个IR简单来说就是可以读取我们的模型并按照模型的顺序搭建出一个可以执行的计算图出来，当然，我们可以对这个计算图进行一系列优化。(现在TVM主推Relay而不是NNVM，Relay可以称为二代NNVM)。

target = 'llvm'

input_name = '0'  # 注意这里为之前导出onnx模型中的模型的输入id，这里为0
shape_dict = {input_name: x.shape}
# 利用Relay中的onnx前端读取我们导出的onnx模型
sym, params = relay.frontend.from_onnx(onnx_model, shape_dict)

网站声明：如果转载，请联系本站管理员。否则一切后果自行承担。

本文链接：https://www.xckfsq.com/news/show.html?id=33059

赞同 0

评论 0 条

利用TVM读取并预测ONNX模型

利用TVM读取并预测ONNX模型

相关文章

关注我们