[PYTHON] TensorFlow 2.X & TensorRT est la première étape pour accélérer l'inférence

introduction

En utilisant le modèle des applications Keras, j'écris brièvement la conversion et l'inférence au modèle pour TensorRT. Les installations telles que TensorRT sont hors de portée car elles utilisent le conteneur NVIDIA GPU Cloud.

Préparation

Conditions préalables

--Docker installé --Et un conteneur GPU est disponible

Démarrez l'environnement d'exécution

Démarrez l'environnement d'exécution avec ce qui suit. Puisque Jupyter est déjà installé, vous pouvez essayer le code suivant sur Jupyter.

docker run -it --rm --gpus all -p 8888:8888 --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 nvcr.io/nvidia/tensorflow:20.02-tf2-py3 bash

# (Optional)
jupyter lab

--nvcr.io / nvidia / tensorflow est un conteneur enregistré dans NVIDIA GPU Cloud. --TensorFlow 2.1, TensorRT 7.0 et d'autres Jupyter sont déjà installés dans ce conteneur.

• (Facultatif) shm-size, ulimit: Comme la mémoire principale est beaucoup utilisée lors de la conversion du modèle, elle est également définie comme contre-mesure en cas d'échec d'allocation de mémoire.
• (Facultatif) Le port 8888 est pour Jupyter.

Référence: Paramètre de croissance

Contraintes de mémoire pour les contre-mesures d'erreur qui se produisent souvent lors de l'utilisation de Keras.

Exemple de limite de mémoire

import tensorflow as tf

for dev in tf.config.experimental.list_physical_devices('GPU'):
    tf.config.experimental.set_memory_growth(dev, True)

conversion

Tout d'abord, le modèle cible est généré. Le point est ** enregistrer le format **. Spécifiez tf pour save_format et enregistrez-le avec le modèle enregistré Tensorflow. (Pas nécessaire car il s'agit de la valeur par défaut dans TensorFlow 2.X)

Stockage de modèles dans Keras

from tensorflow.keras.applications.vgg16 import VGG16

model = VGG16(weights='imagenet')
model.save('./vgg16', save_format='tf')

Transformez ensuite le modèle.

Version virgule flottante simple précision de conversion TensorRT (Float32)

from tensorflow.python.compiler.tensorrt import trt_convert as trt

converter = trt.TrtGraphConverterV2(input_saved_model_dir='./vgg16',
                                    conversion_params=trt.DEFAULT_TRT_CONVERSION_PARAMS)
converter.convert()
converter.save('./vgg16-tensorrt')

Version virgule flottante semi-précision

Si vous souhaitez convertir avec Float16, modifiez les paramètres du convertisseur.

converter = trt.TrtGraphConverterV2(input_saved_model_dir='./vgg16',
                                    conversion_params=trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(precision_mode=trt.TrtPrecisionMode.FP16))

Version entière

Si vous voulez en faire un entier 8 bits, vous devez le calibrer. Je pense que vous devriez utiliser les données utilisées pour apprendre.

Dans ce paramètre VGG16, il est passé sous la forme d'une forme de $ (N, 224, 224, 3) $.

import numpy as np

def calibration_input_fn():  #Fonction de génération de données pour l'étalonnage
    yield np.random.uniform(size=(5, 224, 224, 3)).astype(np.float32),  #À la fin,N'oublie pas

converter = trt.TrtGraphConverterV2(input_saved_model_dir='./vgg16',
                                    conversion_params=trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(precision_mode=trt.TrtPrecisionMode.INT8, use_calibration=True))
converter.convert(calibration_input_fn=calibration_input_fn)
converter.save('./vgg16-tensorrt')

inférence

Chargez le modèle transformé et récupérez les objets utilisés pour l'inférence. Ensuite, l'inférence s'exécute en appelant la fonction de cet objet.

model = tf.saved_model.load('./vgg16-tensorrt', tags=[tf.saved_model.SERVING])
infer = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

#Entrée factice
x = np.random.uniform(size=(3, 224, 224, 3)).astype(np.float32)
#inférence
y = infer(tf.convert_to_tensor(x))['predictions']

Référence: La forme d'entrée peut être prise ci-dessous.

infer.inputs[0].shape
>>> TensorShape([None, 224, 224, 3])

Comparaison des résultats d'exécution

Enfin, une comparaison approximative des résultats d'exécution (TensorRT ci-dessus). Même si cela convient comme ça, la vitesse d'exécution est améliorée. Étant donné que l'utilisation de la mémoire diminuera également, il semble que nous puissions faire diverses choses telles que l'exécution de plusieurs modèles, et je pense que le degré d'avantage changera considérablement en fonction de l'architecture GPU de l'environnement d'exécution.

Environnement d'exécution: GeForce GTX 1080 Ti