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模型的轉換

#Convert PyTorch graph to Relay graph.
import tvm, time
from tvm import relay
input_name = 'input0'  # only one input, set it to this name
shape_list = [(input_name, input_shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)

tvm模型構建

使用導出的torch模型構建tvm模型

#Compile the graph to llvm target with given input specification
target = 'llvm'
target_host = 'llvm'
ctx = tvm.cpu(0)
with relay.build_config(opt_level=3):
    graph, lib, params = relay.build(mod, target=target, target_host=target_host, params=params)
#deploying the compiled model on target.
from tvm.contrib import graph_runtime
tvm_model = graph_runtime.create(graph, lib, ctx)
tvm_model.set_input(**params)

數據準備

隨機生成1000張224*224的圖片用於測試

import numpy as np
#考慮到隨機生成的圖片模型預測結果都一樣，使用*(i % 9 + 1)讓模型預測結果不同
list_img = [np.random.rand(1, 3, 224, 224).astype(np.float32) * (i % 9 + 1) for i in range(1000)]

tvm模型測試

1000張圖片依次預測，結果返回預測類別和累計預測花費時間

def test_tvm(model, list_img):
    since = time.time()
    # 記錄時間
    list_tpoint = []
    # 預測top1的類別
    list_top = []
    for img_x in list_img:
        model.set_input(input_name, tvm.nd.array(img_x))
        # Execute
        model.run()
        # Get outputs
        output = model.get_output(0)
        top1 = np.argmax(output.asnumpy()[0])
        list_top.append(top1)
        list_tpoint.append(time.time() - since)
    return list_top, list_tpoint

torch模型測試

輸出結果同上

@torch.no_grad()
def test_torch(model, list_img):
    since = time.time()
    list_top = []
    list_tpoint = []
    for img_x in list_img:
        img_x = torch.from_numpy(img_x)
        output = model(img_x)
        top1 = np.argmax(output.numpy()[0])
        list_top.append(top1)
        list_tpoint.append(time.time() - since)
    return list_top, list_tpoint

結果分析

tvm_top,tvm_point=test_tvm(tvm_model,list_img)
torch_top,torch_points=test_torch(torch_model,list_img)
#打印2個模型耗費的時間
import matplotlib.pyplot as plt
plt.figure()
plt.plot(tvm_point,label='tvm')
plt.plot(torch_points,label='torch')
plt.legend()
plt.savefig('time_cost.png')
plt.show()

2個模型的部分預測結果

>>> tvm_top[:10]
[111, 111, 318, 818, 818, 644, 644, 644, 644, 111]
>>> torch_top[:10]
[111, 111, 318, 818, 818, 644, 644, 644, 644, 111]

2個模型預測耗時

>>> torch_points[-1]
53.492563009262085
>>> tvm_point[-1]
43.807936906814575

2個模型的累計時間花費圖，橫座標爲圖片數，總座標爲耗時(秒)

小結：TVM相對原生態的torch在推理速度上還是有一定優勢

reference:
https://zhuanlan.zhihu.com/p/88369758
https://docs.tvm.ai/tutorials/frontend/from_pytorch.html