Une approche pour accélérer votre inférence BERT avec ONNX-TorchScript

Ces dernières années, les modèles basés sur l'architecture Transformer ont été à l'origine des avancées de la PNL dans la recherche et l'industrie. BERT, XLNET, GPT ou XLM sont quelques-uns des modèles qui ont amélioré l'état de l'art et ont atteint le sommet des benchmarks populaires tels que GLUE.

Ces avancées s'accompagnent d'un coût de calcul élevé, la plupart des modèles basés sur des transformateurs sont massifs et le nombre de paramètres et de données utilisés pour la formation ne cessent d'augmenter. Alors que le modèle BERT original comportait déjà 110 millions de paramètres, le dernier GPT-3 en compte 175 milliards, un chiffre impressionnant ~1700x augmentation de la valeur de la recherche en deux ans.

Ces modèles massifs nécessitent généralement des centaines de GPU pendant plusieurs jours d'entraînement pour être efficaces. Heureusement, grâce à l'apprentissage par transfert, nous pouvons télécharger des modèles préentraînés et les affiner rapidement sur nos propres ensembles de données beaucoup plus petits à moindre coût.

Cela dit, une fois la formation terminée, vous avez encore un énorme modèle à portée de main que vous souhaiterez peut-être déployer en production. L'inférence prend un temps relativement long par rapport aux modèles plus modestes et elle peut être trop lente pour atteindre le débit dont vous avez besoin.

Bien que vous puissiez investir dans du matériel plus rapide ou utiliser davantage de serveurs pour effectuer le travail, il existe différentes manières de réduire le temps d'inférence de votre modèle :

• Modèle d'élagage : Réduisez le nombre de couches, la dimension des intégrations ou le nombre d'unités dans les couches masquées.
• Quantification : Au lieu d'utiliser un nombre flottant de 32 bits (FP32) pour les pondérations, utilisez une demi-précision (FP16) ou même un entier de 8 bits.
• Exporter un modèle depuis Pytorch/Tensorflow natif vers un format ou un moteur d'inférence approprié (TorchScript/ONNX/Tensorrt...)
• Traitement par lots: Prédisez sur des lots d'échantillons plutôt que sur des échantillons individuels

La première et la deuxième approche impliquent généralement une reconversion de votre modèle, tandis que les deux dernières approches sont effectuées après l'entraînement et sont essentiellement indépendantes de votre tâche particulière.

Si la vitesse d'inférence est extrêmement importante pour votre cas d'utilisation, vous devrez probablement expérimenter toutes ces méthodes pour produire un modèle fiable et incroyablement rapide. Dans la plupart des cas, cependant, l'exportation de votre modèle vers un format/framework approprié et la prédiction par lots vous donneront des résultats beaucoup plus rapides pour un minimum de travail. Nous nous concentrerons ici sur cette approche pour voir l'impact qu'elle peut avoir sur débit de notre modèle.

Nous explorerons les effets de la modification du format du modèle et du traitement par lots à l'aide de quelques expériences :

• Base de référence avec Pytorch à la vanille CPU/GPU
• Exporter le modèle Pytorch vers Script Torch CPU/GPU
• Modèle Pytorch pour ONNX CPU/GPU
• Toutes les expériences ont été réalisées sur des lots d'échantillons 1/2/4/8/16/32/64

À partir de cet article, il n'est pas encore possible d'exporter directement un modèle de transformateur de Pytorch vers TensorRT en raison du manque de support de int64 utilisé par les intégrations Pytorch. Nous allons donc l'ignorer pour le moment.

Nous allons effectuer une classification des phrases sur Camembert (~100M de paramètres), une variante française de Roberta. Étant donné que la grande majorité des calculs sont effectués dans le modèle du transformateur, vous devriez obtenir des résultats similaires quelle que soit votre tâche.

Nous allons d'abord examiner rapidement comment exporter un modèle Pytorch vers le format/framework approprié. Si vous ne voulez pas lire le code, vous pouvez passer à la section des résultats plus bas.

Comment exporter votre modèle

Pytorch à la vanille

Enregistrer et charger un modèle dans Pytorch est assez simple, bien qu'il existe différentes manières de procéder. Par déduction, la documentation officielle recommande de sauvegarder le « state_dict » de votre modèle, qui est un dictionnaire python contenant les paramètres apprenables de votre modèle. C'est plus léger et plus robuste que le décapage de l'ensemble de votre modèle.

#savingmodel = SequenceClassifier () train_model (model) torch.save (model.state_dict (), 'pytorch_model.pt') #loadingmodel = SequenceClassifier () model.load_state_dict (torch.load (PATH)) model.eval () #Set couches de suppression et de normalisation par lots vers le mode évaluation avec torch.go_grad () : logits = model (**lot_x)

Suivez ceci lien pour plus d'informations sur l'enregistrement/le chargement sur Pytorch.

Torchscript JIT

TorchScript est un moyen de créer des modèles sérialisables et optimisables à partir de votre code Pytorch. Une fois exporté vers Torchscript, votre modèle pourra être exécuté depuis Python et C++.

• Trace: une entrée est envoyée via le modèle et toutes les opérations sont enregistrées dans un graphique qui définira votre modèle Torchscript.
• Scénario: Si votre modèle est plus complexe et possède un flux de contrôle tel que des instructions conditionnelles, les scripts inspecteront le code source du modèle et le compileront en tant que code TorchScript.

Notez que puisque votre modèle sera sérialisé, vous ne pourrez pas le modifier une fois qu'il aura été enregistré. Vous devez donc le mettre en mode évaluation et l'exporter sur l'appareil approprié avant de l'enregistrer.

Si vous voulez faire des inférences à la fois sur le processeur et le GPU, vous devez enregistrer 2 modèles différents.

#savingjit_sample = (batch_x ['input_ids'] .int () .to (périphérique), batch_x ['attention_mask'] .int () .to (appareil)) model.eval () model.to (périphérique) module = torch.jit.trace (modèle, jit_sample) torch.jit.save ('model_jit.pt') #loadingmodel = torch.jit..be load ('model_jit.pt', map_location=torch.device (device)) logits = model (**batch_x)

Pour une introduction plus complète, vous pouvez suivre le didacticiel.

ONNX

ONNX fournit un format open source pour les modèles d'IA, la plupart des frameworks peuvent exporter leur modèle au format ONNX. Outre l'interopérabilité entre les frameworks, ONNX s'accompagne d'une certaine optimisation qui devrait accélérer l'inférence.

L'exportation vers ONNX est un peu plus compliquée, mais Pytorch fournit une fonction d'exportation directe, vous n'avez qu'à fournir quelques informations clés.

• opset_version, pour chaque version, un ensemble d'opérateurs sont pris en charge, certains modèles aux architectures plus exotiques ne sont peut-être pas encore exportables.
• noms_entrées et noms_sorties sont les noms à attribuer aux nœuds d'entrée et de sortie du graphe.
• axes_dynamiques argument est un dictionnaire qui indique quelle dimension de vos variables d'entrée et de sortie peut changer, par exemple le batch_size ou la longueur de la séquence.

‍

#savinginput_x = jit_sample ## taking sample from previous exampletorch.onnx.export(model, input_x,'model_onnx.pt',export_params=True,  opset_version=11, do_constant_folding=True, input_names = ['input_ids', 'attention_mask'], output_names = ['output'],dynamic_axes= {'input_ids' : {0 : 'batch_size', 1:'length'},'attention_mask' : {0 : 'batch_size', 1:'length'},'output' : {0 : 'batch_size'}})#loadingmodel = onnxruntime.InferenceSession(model_onnx)batch_x = {'input_ids':sample['input_ids'].cpu().numpy(),"attention_mask":sample['attention_mask'].cpu().numpy()}logits = model.run(None, batch_x)

Le runtime ONNX peut être utilisé avec un GPU, bien qu'il nécessite des versions spécifiques de CUDA, cuDNN et OS, ce qui rend le processus d'installation difficile au début.

Pour un tutoriel plus complet, vous pouvez suivre le documentation.

Résultats expérimentaux

Chaque configuration a été exécutée 5 fois sur un ensemble de données de 1 000 phrases de différentes longueurs. Nous avons testé 2 GPU populaires différents : T4 et V100 avec Torch 1.7.1 et ONNX 1.6.0. N'oubliez pas que les résultats peuvent varier en fonction de votre matériel, de la version des packages et de l'ensemble de données.

Le temps d'inférence varie d'environ 50 ms par échantillon en moyenne à 0,6 ms sur notre ensemble de données, en fonction de la configuration matérielle.

Sur le processeur, le format ONNX est clairement gagnant pour batch_size <32, auquel cas le format ne semble plus vraiment avoir d'importance. Si nous prédisons échantillon par échantillon, nous constatons que ONNX parvient à être aussi rapide que l'inférence sur notre base de référence sur le GPU pour une fraction du coût.

Comme prévu, l'inférence est beaucoup plus rapide sur un GPU, en particulier avec une taille de lot plus élevée. Nous pouvons également constater que la taille de lot idéale dépend du GPU utilisé :

• Pour le T4, la meilleure configuration est d'exécuter ONNX avec des lots de 8 échantillons, ce qui donne un ~12x accélération par rapport à la taille de lot 1 sur Pytorch
• Pour le V100 avec des lots de 32 ou 64, nous pouvons obtenir jusqu'à ~28x accélération par rapport à la base de référence pour le GPU et ~90x pour la référence sur le processeur.

Dans l'ensemble, nous constatons que le choix d'un format approprié a un impact significatif pour les lots de plus petite taille, mais que cet impact diminue à mesure que les lots grossissent. Pour les lots de 64 échantillons, les 3 configurations se situent à environ 10 % l'une de l'autre.

Impact de la longueur des séquences et de la stratégie de traitement par lots

Une autre chose à prendre en compte est la longueur de la séquence. Les transformateurs sont généralement limités à des séquences de 512 jetons, mais il existe une énorme différence en termes de vitesse et de mémoire requise pour les différentes longueurs de séquences dans cette plage.

Le temps d'inférence augmente à peu près linéairement avec la longueur de la séquence pour les lots plus importants, mais pas pour les échantillons individuels. Cela signifie que si vos données sont constituées de longues séquences de texte (articles de presse par exemple), vous n'obtiendrez pas une accélération aussi importante en cas de traitement par lots. Comme toujours, cela dépend de votre matériel, un V100 est plus rapide qu'un T4 et ne souffrira pas autant lors de la prédiction de longues séquences, alors que d'un autre côté, notre processeur est complètement dépassé :

Si vos données sont hétérogènes sur le plan de la longueur et que vous travaillez avec des lots, ces divergences poseront des problèmes en raison de la nécessité de placer vos échantillons sur le plus long de votre lot, ce qui demande beaucoup de calculs. Par conséquent, il est généralement préférable de regrouper des échantillons de longueur similaire, car il est probablement plus rapide de prévoir plusieurs lots de longueur similaire qu'un seul gros lot composé principalement de jetons de remplissage.

En guise de vérification rapide, voyons ce qui se passe lorsque nous trions notre ensemble de données avant d'exécuter l'inférence :

Comme nous nous y attendions pour des lots de plus grande taille, il existe une incitation significative à regrouper des échantillons de longueur similaire. Pour les données non triées, à mesure que les lots grossissent, il y a une probabilité croissante de se retrouver avec des échantillons plus longs, ce qui augmentera considérablement le temps d'inférence de l'ensemble du lot. Nous pouvons constater que le fait de passer de 16 à 64 batch_size ralentit l'inférence de 20 % alors qu'elle est 10 % plus rapide avec les données triées.

Cette stratégie peut également être utilisée pour réduire considérablement votre temps d'entraînement, mais cela doit être fait avec prudence car elle peut avoir un impact négatif sur les performances de votre modèle, en particulier s'il existe une certaine corrélation entre vos étiquettes et la longueur de vos échantillons.

Prochaines étapes

Bien que ces expériences aient été exécutées directement en Python, les modèles Torchscript et ONNX peuvent être chargés directement en C++, ce qui pourrait améliorer encore la vitesse d'inférence.

Si votre modèle est encore trop lent pour votre cas d'utilisation, Pytorch propose différentes options de quantification. La « quantification dynamique » peut être effectuée après l'entraînement, mais elle aura probablement un impact sur la précision de votre modèle, tandis que la « formation tenant compte de la quantification » nécessite un nouvel entraînement mais devrait avoir un impact moindre sur les performances de votre modèle.

Conclusion

Comme nous l'avons vu, il n'existe pas de solution simple pour optimiser votre temps d'inférence, car cela dépend principalement de votre matériel spécifique et du problème que vous essayez de résoudre. Par conséquent, vous devez effectuer vos expériences avec votre propre matériel cible et vos propres données pour obtenir des résultats fiables.

Néanmoins, certaines directives devraient être vraies et faciles à mettre en œuvre :

• La prédiction des lots peut permettre d'accélérer considérablement jusqu'à une certaine taille (en fonction de votre matériel spécifique), en particulier si vous pouvez regrouper des échantillons de longueur similaire
• L'utilisation de Torchscript ou ONNX permet une accélération significative pour réduire la taille des lots et la longueur des séquences, l'effet est particulièrement fort lors de l'exécution d'une inférence sur des échantillons individuels.
• ONNX semble être la plus performante des trois configurations que nous avons testées, bien qu'elle soit également la plus difficile à installer pour l'inférence sur le GPU
• Torchscript fournit une accélération fiable pour les lots de petite taille et est très facile à configurer.