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
那么仿照上面例子的写法,就可以开始写自己的测试代码.大体上一样,只是在输入的数据用随机生成的.
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的推理速度
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'](ctx))
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的核心代码很简单
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),得到调优结果之后再次利用调优结果编译测试推理速度
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"](dev))
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)
同时可以参考大佬给出的例子github.com/AndrewZhaoL…