theme: smartblue

TVM VS TensorRT

上一次記錄瞭如何安裝TVM,同時也説了會將TVM與TensorRT各自優化後的模型做一下性能比較.這一次就選擇resnet18來比較一下不同框架下的推理速度

1. 準備

首先確保有TensorRT以及TVM環境,然後可以去onnx下載我們測試的模型文件,這裏我選擇的是resnet18-v2-7

然後轉為trt文件,trtexec --onnx=resnet18-v2-7.onnx --saveEngine=resnet18v2.trt

這裏不選擇--fp16,因為一開始在TVM中沒有使用fp16導致後面tune完之後不想改了(tune花費的時間超乎想象!!!)為了公平起見TensorRT也不用fp16.

2. TensorRT

使用python的TensorRT推理,具體步驟可以先看看官網給出的例子,當然TensorRT編譯的目錄下也有相應的examples可以參考,這裏就給出官方的例子Using PyTorch with TensorRT through ONNX

參考其中的代碼,發現TensorRT推理主要步驟如下:

載入engine/trt,並且反序列化
構建運行器
生成輸入輸出數據並且創建GPU中的緩存空間
將輸入數據從CPU轉到GPU進行推理,然後將結果從GPU放回CPU

那麼仿照上面例子的寫法,就可以開始寫自己的測試代碼.大體上一樣,只是在輸入的數據用隨機生成的.

```python import numpy as np import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit import time import torch

Batch_size

BATCH_SIZE=1

讀取trt文件

f=open("./resnet18v2.trt","rb")

定義輸入輸出name

input_name="data" output_name="resnetv22_dense0_fwd"

轉為trt時是否使用了--fp16

USE_FP16=False target_type=np.float16 if USE_FP16 else np.float32

開始反序列化推理

logger = trt.Logger(trt.Logger.WARNING) runtime=trt.Runtime(logger)

反序列化得到engine

engine=runtime.deserialize_cuda_engine(f.read()) f.close()

創建執行上下文對象

context = engine.create_execution_context() input_idx = engine[input_name] output_idx = engine[output_name]

在GPU中創建緩衝區

buffers=[None]*2

在cpu中設置輸入輸出大小,類型的buffer,在GPU中創建buffer

input_data=cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)),dtype=target_type) output_data=cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)),dtype=target_type)

d_input=cuda.mem_alloc(1 * input_data.nbytes) d_output = cuda.mem_alloc(1 * output_data.nbytes)

buffers[input_idx]=d_input buffers[output_idx]=d_output

stream=cuda.Stream()

推理

def predict(datas): # 將輸入數據從CPU轉移到GPU cuda.memcpy_htod_async(buffers[input_idx],datas,stream) stream.synchronize() start=time.time() context.execute_async_v2(buffers,stream.handle) stream.synchronize() cost_time=time.time()-start # 將推理結果從GPU下載到CPU cuda.memcpy_dtoh_async(output_data,buffers[output_idx],stream) print(f"predict label:{output_data.argmax()},cost time:{cost_time*1000}ms") return cost_time

創建測試數據

datas= np.random.randn(BATCH_SIZE, 224, 224, 3).astype(target_type)

對GPU Warm up再進行測試

循環查看每次推理時間,並計算平均時間

Test_times=20 Warm_times=10

time_sum=0 for i in range(Warm_times+Test_times): if i>=Warm_times: if i==Warm_times: print("End Warming up------------------------------------") time_sum+=predict(datas) else: if i==0: print("Warming up...") predict(datas) print(f"avg cost time:{time_sum/Test_times*1000}ms") ```

可以看到TensorRT花費的時間大概是0.32ms左右

3. TVM

這次輪到TVM了,首先看看官網給出的教程Quick Start Tutorial for Compiling Deep Learning Models

大致分為下面幾步:

在Relay中定義網絡,支持包括onnx的很多數據模型,這裏直接用上面的onnx就好
編譯模型得到lib
通過lib生成Graph,並且執行Graph

3.1 Baseline

按照教程,同樣可以寫一個簡單的demo得到一個Baseline的推理速度

```python import tvm from tvm import te from tvm import relay from tvm.contrib import graph_executor

import numpy as np import torch import time import onnx

BATCH_SIZE=1 dtype=np.float32

input_shape = [1, 3, 224, 224] input_data = torch.randn(input_shape) output_data=np.empty((BATCH_SIZE,1000)).astype(dtype)

onnx_model= onnx.load("./resnet18-v2-7.onnx")

input_name='data' shape_dict={input_name:input_shape}

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

target=tvm.target.cuda()

ctx=tvm.cuda()

with tvm.transform.PassContext(opt_level=opt_level): lib=relay.build(mod,target='cuda -libs=cudnn,cublas',params=params)

model=graph_executor.GraphModule(lib'default')

def pred(): model.set_input(input_name,input_data) torch.cuda.synchronize() start=time.time() model.run() torch.cuda.synchronize() cost=time.time()-start output=model.get_output(0).asnumpy()

print(f"cost time:{cost*1000}ms")
return cost

print("Untuned testing...") Test_times=20 Warm_times=10 cost_time=0 for i in range(Warm_times+Test_times): if i>=Warm_times: cost_time+=pred() else: pred() print(f"avg cost time:{cost_time/Test_times*1000}ms") ```

這裏編譯使用的target一般如果是GPU的話就是cuda,但是網上搜索看到這樣寫可以加入cudnn等優化.

Baseline的推理時間是0.76ms左右,根據上面的輸出也可以看到有些算子沒有被tune,那為了追求性能,就來tune一下試試

3.2 Tune

這裏使用TVMC來對模型進行tune,TVMC是一個封裝好的高級API可以非常簡單實現調優過程.具體例子可以看Getting Starting using TVMC Python: a high-level API for TVM

例子中給出的是以llvm也就是CPU優化,但是我們需要的是GPU上的優化,所以得稍微修改一下,其中tune的核心代碼很簡單

```python from tvm.driver import tvmc log_file = "tune.json" model=tvmc.load("./resnet18-v2-7.onnx",shape_dict={'data':[BATCH_SIZE,3,224,224]}) package = tvmc.compile(model, target="cuda")

tvmc.tune(model, target="cuda",tuning_records=log_file) ```

還是一樣先導入onnx模型,然後用tvmc編譯再tune,這裏一定要將tune的結果保存下來為了以後重用,不然tune一次會花費很長時間. 在這裏插入圖片描述

從輸出體會一下吧,tune一次大概得花費7小時左右(第一次忘記保存優化結果,跑了兩次tune),得到調優結果之後再次利用調優結果編譯測試推理速度

```python from tvm import autotvm log_file='tune.json' mod, params = relay.frontend.from_onnx(onnx_model,shape_dict) with autotvm.apply_history_best(log_file): with tvm.transform.PassContext(opt_level=4, config={}): lib = relay.build(mod, target="cuda -libs=cudnn,cublas", params=params) dev=tvm.device("cuda")

model=graph_executor.GraphModule(lib"default")

print("tuned testing...") Test_times=20 Warm_times=10 cost_time=0 for i in range(Warm_times+Test_times): if i>=Warm_times: cost_time+=pred() else: pred() print(f"avg cost time:{cost_time/Test_times*1000}ms") ```

可以看出tune之後從0.76ms降到了0.56ms左右

最後

上個結果對比

| 框架 | 平均推理時間(ms) | 最大推理時間(ms) | 最小推理時間(ms) | | -------- | ---------------- | ---------------- | ---------------- | | TensorRT | 0.3195 | 0.6337 | 0.3023 | | TVM Base | 0.7672 | 1.6467 | 0.6861 | | TVM Tune | 0.5652 | 1.0623 | 0.5562 |

費了九牛二虎之力調優出來的模型還是比不過TensorRT的推理速度,七小時調優只減少0.2ms的推理時間,但起碼模型的穩定性得到了提升還是有些許安慰的.

當然除了用tvmc進行調優,也可以用低級的python接口AutoTVM進行調優,可以參考官方文檔Compiling and Optimizing a Model with the Python Interface (AutoTVM)

如果需要使用TVM進行fp16推理,可以在build lib之前對mod進行混合精度

mod=tvm.relay.transform.ToMixedPrecision(mixed_precision_type='float16')(mod)

同時可以參考大佬給出的例子http://github.com/AndrewZhaoLuo/CenterFaceTVMDemo/blob/main/scripts/example_tvm_tune.py

TVM VS TensorRT推理速度比較