在网上关于TensorRT的资料很杂乱,这里是我自己在GTX1070和GTX1080ti上做测试的结果。包括tensorflow(keras)和tensorrt性能对比,TensorRT如何多线程运行。模型转化的方法。
ubuntu 16.04
TensorRT 5.0.2.6
python 3.5
cuda 9.0
cudnn 7.3.0
使用keras版本和使用tensorrt加速后的FPS对比如下表。
GTX1070
模型类别 | keras(tensorflow) | TensorRT |
---|---|---|
vgg16 | 58 | 87 |
yolov3 | 13 | 18 |
GTX1080ti
模型类别 | keras(tensorflow) | TensorRT |
---|---|---|
vgg16 | 95 | 136 |
yolov3 | 17 | 23 |
性能上面提高的不是很多,我后来测试发现不管是使用tensorflow还是tensorrt,他们在推理上面消耗的时间基本是差不多的,唯一的区别就是在处理数据上面需要的时间不同,下面是我在GTX1070上面用yolov3的消耗的时间表,可以看出大部分消耗的时间是在数据处理上面,我们可以在数据处理上做进一步的优化,以提高速度。
类型 | 耗时(ms) |
---|---|
准备数据 | 6 |
推理 | 18 |
处理结果 (将矩阵转化为class name) |
26 |
画图 | 1 |
总时间 | 51 |
python yolov3_to_onnx.py
python onnx_to_tensorrt.py
python multi_thread_detect.py --mode yolov3
关键代码是 31行到43行的代码
cuda_ctx = cuda.Device(0).make_context() # important
""" 创建模型类的方法必须要在线程的run方法中 """
if self.mode == "yolov3":
model = Yolov3TRT()
elif self.mode == "vgg16":
model = VGG16(overall_model_path="vgg16.uff")
matrix = self.queue.get(block=True)
# print detection info
model.predict(matrix)
cuda_ctx.pop() # important
del cuda_ctx # important
请克隆此处代码,使用下面的命令进行模型的转化
python keras_to_tensorflow.py
--input_model="path/to/keras/model.h5"
--output_model="path/to/save/model.pb"
convert-to-uff [your_pb_model]
转化之后会在终端打印出uff模型的input节点和output节点的名称,要运行除vgg16以外的模型的时候只需要在创建Classifier的时候输入所需的model_input_name、model_output_name、input_shape和labels即可。
python multi_thread_detect.py --mode vgg16
2、安装TensorRT的方法,这里使用的是tar包安装