TVM笔记一

import torch
import numpy as np
import torchvision
import onnx
import tvm
from tvm import te
import tvm.relay as relay

# 加载torchvision中ResNet18模型
model_name = "resnet18"
model = getattr(torchvision.models, model_name)(pretrained=True)
model = model.eval()  # 不启用BatchNormization和DropOut 测试的时候用

from PIL import Image
image_path = 'cat.png'
img = Image.open(image_path).resize((224, 224))

# 处理图像并转成Tensor
from torchvision import transforms
my_preprocess = transforms.Compose(
    [
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)
img = my_preprocess(img)
img = np.expand_dims(img, 0)

# 使用Pytorch进行预测结果
with torch.no_grad():
    torch_img = torch.from_numpy(img)
    output = model(torch_img)
    top1_torch = np.argmax(output.numpy())
print(top1_torch)

# 模型存储为 onnx
torch_out = torch.onnx.export(model, torch_img, 'resnet18_1.onnx', verbose=True, export_params=True)

# 加载 onnx 模型
onnx_model = onnx.load('resnet18_1.onnx')

# 设置 target 为 llvm 表示要在 CPU 后端运行Relay IR
target = "llvm"
input_name = "input.1"
shape_dict = {input_name: img.shape}
mod, params = relay.frontend.from_onnx(onnx_model, shape_dict)

# RelayIR 编译为 RuntimeModel
with tvm.transform.PassContext(opt_level=1):
    intrap = relay.build_module.create_executor("graph", mod, tvm.cpu(0), target)
    print(intrap)

# execute on tvm
dtype = "float32"
tvm_output = intrap.evaluate()(tvm.nd.array(img.astype(dtype)), **params).asnumpy()

print(np.argmax(tvm_output))

# 另一种构建build
# 设置优化级别
with tvm.transform.PassContext(opt_level=3):
    #编译模型
    lib = relay.build(mod, target, params=params)
    print(lib)