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推理主要步骤如下:

1. 载入engine/trt,并且反序列化
2. 构建运行器
3. 生成输入输出数据并且创建GPU中的缓存空间
4. 将输入数据从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

大致分为下面几步:

1. 在Relay中定义网络,支持包括onnx的很多数据模型,这里直接用上面的onnx就好
2. 编译模型得到lib
3. 通过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…