pytorch笔记：12)TVM-Pytorch模型编译体验+性能测试

tvm简介：https://zhuanlan.zhihu.com/p/88369758

实验环境

18.04.1-Ubuntu x86_64 x86_64 x86_64 GNU/Linux
CPU: 8  Intel(R) Core(TM) i7-6700 CPU @ 3.40GHz
tvm-0.7.dev1
torch-1.4.0
torchvision-0.5.0

tvm安装文档，另官网推荐使用torch不小于1.3.0

torch模型导出

使用torchvision自带的resnet18进行实验测试

import torch
import torchvision

model_name = 'resnet18'
model = getattr(torchvision.models, model_name)(pretrained=True)
# ~/.cache/torch/checkpoints/resnet18-5c106cde.pth
model = model.eval()

# We grab the TorchScripted model via tracing
input_shape = [1, 3, 224, 224]
input_data = torch.randn(input_shape)
#scripted_model不依赖python环境
scripted_model = torch.jit.trace(model, input_data).eval()

最后一句会提示警告，论坛帖子说这是正常的- -

WARNING:root:Untyped Tensor found, assume it is float

模型转换

torch模型到tvm模型的转换