利用TensorRT转换ResNet50

2022.08.16

本文使用Tensorflow中集成的TensorRT进行模型转换。

不需要安装TensorRT的Python库，但是TensorRT还需要安装，需要用到的包是libnvinfer，如若不安装TensorRT，需要使用apt-get install libnvinfer进行安装。

关于TensorRT的安装，请参考：

为Tensorflow安装TensorRT（tar）_January_Cao的博客-CSDN博客

0. 文件夹构成

├── compare_models.py
├── convert.py
├── data

│   ├── download_images.sh
│   ├── img0.jpg
│   ├── img1.jpg
│   ├── img2.jpg
│   └── img3.jpg

├── model_check.py
└── save_model.py

1. 准备数据

也可以随意从网上下载图片。

或者执行【sh download_images.sh】下载图片至data文件夹

wget -O img0.jpg https://res.cloudinary.com/roundglass/image/upload/q_auto/ar_16:9,c_fill,w_1250/f_auto,g_auto/v1639561104/rg/collective/media/vyy8zc9ww3gqsjkqy9gk.jpgwget -O img1.jpg https://www.fundacion-affinity.org/sites/default/files/los-10-sonidos-principales-del-gato-fa.jpgwget -O img2.jpg https://hips.hearstapps.com/hmg-prod.s3.amazonaws.com/images/dog-puppy-on-garden-royalty-free-image-1586966191.jpg?crop=1.00xw:0.669xh;0,0.190xh&resize=1200:*wget -O img3.jpg https://upload.wikimedia.org/wikipedia/commons/thumb/0/05/Parrot.jpg/1280px-Parrot.jpg

2. 模型数据下载和推理，保存模型为saved_model

执行以下代码，可以看到对图片的推理结果，并把模型以saved_model的形式保存到‘resnet50_saved_model’路径下。

# save_model.py
import numpy as npfrom tensorflow.keras.applications.resnet50 import ResNet50
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet50 import preprocess_input, decode_predictions# load model
model = ResNet50(weights='imagenet')for i in range(4):img_path = './data/img{}.jpg'.format(i)img = image.load_img(img_path, target_size=(224, 224))x = image.img_to_array(img)x = np.expand_dims(x, axis=0)x = preprocess_input(x)# predictionpreds = model.predict(x)print('{} - Predicted: {}'.format(img_path, decode_predictions(preds, top=3)[0]))# save model to resnet50_saved_model path
model.save('resnet50_saved_model')

3. 查看结果

1/1 [==============================] - 2s 2s/step
./data/img0.jpg - Predicted: [('n01742172', 'boa_constrictor', 0.8889929), ('n01729322', 'hognose_snake', 0.04273584), ('n01734418', 'king_snake', 0.032792028)]
1/1 [==============================] - 0s 14ms/step
./data/img1.jpg - Predicted: [('n02124075', 'Egyptian_cat', 0.53852856), ('n02127052', 'lynx', 0.19898652), ('n02123159', 'tiger_cat', 0.1614549)]
1/1 [==============================] - 0s 14ms/step
./data/img2.jpg - Predicted: [('n02113799', 'standard_poodle', 0.5422854), ('n02093647', 'Bedlington_terrier', 0.44937962), ('n02091134', 'whippet', 0.002195063)]
1/1 [==============================] - 0s 14ms/step
./data/img3.jpg - Predicted: [('n01818515', 'macaw', 0.7934733), ('n01820546', 'lorikeet', 0.11080903), ('n01828970', 'bee_eater', 0.08086483)]

4. 利用保存的saved_model推测结果

导入刚刚保存的‘resnet50_saved_model’，结果和之前一致。

# model_check.py
import numpy as np
import tensorflow as tffrom tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet50 import preprocess_input, decode_predictions# load model
model = tf.keras.models.load_model('resnet50_saved_model')for i in range(4):img_path = './data/img{}.jpg'.format(i)img = image.load_img(img_path, target_size=(224, 224))x = image.img_to_array(img)x = np.expand_dims(x, axis=0)x = preprocess_input(x)# predictionpreds = model.predict(x)print('{} - Predicted: {}'.format(img_path, decode_predictions(preds, top=3)[0]))

1/1 [==============================] - 2s 2s/step
./data/img0.jpg - Predicted: [('n01742172', 'boa_constrictor', 0.8889929), ('n01729322', 'hognose_snake', 0.04273584), ('n01734418', 'king_snake', 0.032792028)]
1/1 [==============================] - 0s 20ms/step
./data/img1.jpg - Predicted: [('n02124075', 'Egyptian_cat', 0.53852856), ('n02127052', 'lynx', 0.19898652), ('n02123159', 'tiger_cat', 0.1614549)]
1/1 [==============================] - 0s 15ms/step
./data/img2.jpg - Predicted: [('n02113799', 'standard_poodle', 0.5422854), ('n02093647', 'Bedlington_terrier', 0.44937962), ('n02091134', 'whippet', 0.002195063)]
1/1 [==============================] - 0s 16ms/step
./data/img3.jpg - Predicted: [('n01818515', 'macaw', 0.7934733), ('n01820546', 'lorikeet', 0.11080903), ('n01828970', 'bee_eater', 0.08086483)]

5. 转换为精度为Float32模型

执行以下代码把刚刚的‘resnet50_saved_model’利用TensorRT优化模型，并把优化好的模型保存在‘resnet50_saved_model_TFTRT_FP32’

# convert.py
from tensorflow.python.compiler.tensorrt import trt_convert as trtconversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(precision_mode=trt.TrtPrecisionMode.FP32)converter = trt.TrtGraphConverterV2(input_saved_model_dir='resnet50_saved_model',conversion_params=conversion_params)
converter.convert()
converter.save(output_saved_model_dir='resnet50_saved_model_TFTRT_FP32')

6. 运行时间对比

比较保存的saved_model‘resnet50_saved_model’和用TensorRT优化后的模型‘resnet50_saved_model_TFTRT_FP32’的Throughput(每秒钟处理的图像数)，结果如下，优化后的模型是之前的1.7倍。

	resnet50_saved_model	resnet50_saved_model_TFTRT_FP32
Throughput[images/s]	340	577

# compare_models.py
import timeimport numpy as npimport tensorflow as tf
from tensorflow.python.saved_model import tag_constants, signature_constants
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet50 import preprocess_inputmodels = ['resnet50_saved_model', 'resnet50_saved_model_TFTRT_FP32']for model in models:saved_model = tf.keras.models.load_model(model)infer = saved_model.signatures['serving_default']batch_size = 8batched_input = np.zeros((batch_size, 224, 224, 3), dtype=np.float32)for i in range(batch_size):img_path = 'data/img%d.jpg' % (i % 4)img = image.load_img(img_path, target_size=(224, 224))x = image.img_to_array(img)x = np.expand_dims(x, axis=0)x = preprocess_input(x)batched_input[i, :] = xbatched_input = tf.constant(batched_input)N_warmup_run = 50N_run = 1000elapsed_time = []for i in range(N_warmup_run):preds = infer(batched_input)for i in range(N_run):start_time = time.time()preds = infer(batched_input)end_time = time.time()elapsed_time = np.append(elapsed_time, end_time - start_time)print('Throughput: {:.0f} images/s'.format(N_run * batch_size / elapsed_time.sum()))

如果把精度改为FLOAT16之后，处理速度会更快。