TVM VS TensorRT推理速度比較
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