torch2trt
v0.5.0
Quels modèles utilisez-vous ou espérez-vous utiliser avec Tensorrt? N'hésitez pas à rejoindre la discussion ici.
Torch2trt est un convertisseur pytorch à Tensorrt qui utilise l'API Tensorrt Python. Le convertisseur est
Facile à utiliser - Convertir les modules avec une seule fonction d'appel torch2trt
Facile à étendre - Écrivez votre propre convertisseur de couche dans Python et enregistrez-le avec @tensorrt_converter
Si vous trouvez un problème, faites-le nous savoir!
Veuillez noter que ce convertisseur a une couverture limitée de Tensorrt / Pytorch. Nous l'avons créé principalement pour optimiser facilement les modèles utilisés dans le projet Jetbot. Si vous trouvez le convertisseur utile avec d'autres modèles, veuillez nous le faire savoir.
Vous trouverez ci-dessous quelques exemples d'utilisation, pour plus, consultez les ordinateurs portables.
import torch
from torch2trt import torch2trt
from torchvision . models . alexnet import alexnet
# create some regular pytorch model...
model = alexnet ( pretrained = True ). eval (). cuda ()
# create example data
x = torch . ones (( 1 , 3 , 224 , 224 )). cuda ()
# convert to TensorRT feeding sample data as input
model_trt = torch2trt ( model , [ x ]) Nous pouvons exécuter le TRTModule retourné comme le modèle Pytorch d'origine
y = model ( x )
y_trt = model_trt ( x )
# check the output against PyTorch
print ( torch . max ( torch . abs ( y - y_trt ))) Nous pouvons enregistrer le modèle en tant que state_dict .
torch . save ( model_trt . state_dict (), 'alexnet_trt.pth' ) Nous pouvons charger le modèle enregistré dans un TRTModule
from torch2trt import TRTModule
model_trt = TRTModule ()
model_trt . load_state_dict ( torch . load ( 'alexnet_trt.pth' ))Nous avons testé le convertisseur contre ces modèles à l'aide du script Test.sh. Vous pouvez générer les résultats en appelant
./test.sh TEST_OUTPUT.mdLes résultats ci-dessous montrent le débit en FPS. Vous pouvez trouver la sortie brute, qui comprend la latence, dans le dossier Benchmarks.
| Modèle | Nano (pytorch) | Nano (Tensorrt) | Xavier (pytorch) | Xavier (Tensorrt) |
|---|---|---|---|---|
| Alexnet | 46.4 | 69.9 | 250 | 580 |
| Squeezenet1_0 | 44 | 137 | 130 | 890 |
| Squeezenet1_1 | 76.6 | 248 | 132 | 1390 |
| resnet18 | 29.4 | 90.2 | 140 | 712 |
| resnet34 | 15.5 | 50.7 | 79.2 | 393 |
| resnet50 | 12.4 | 34.2 | 55,5 | 312 |
| resnet101 | 7.18 | 19.9 | 28.5 | 170 |
| resnet152 | 4.96 | 14.1 | 18.9 | 121 |
| densenet121 | 11.5 | 41.9 | 23.0 | 168 |
| densenet169 | 8.25 | 33.2 | 16.3 | 118 |
| densenet201 | 6.84 | 25.4 | 13.3 | 90.9 |
| densenet161 | 4.71 | 15.6 | 17.2 | 82.4 |
| VGG11 | 8.9 | 18.3 | 85.2 | 201 |
| VGG13 | 6.53 | 14.7 | 71.9 | 166 |
| VGG16 | 5.09 | 11.9 | 61.7 | 139 |
| VGG19 | 54.1 | 121 | ||
| VGG11_BN | 8.74 | 18.4 | 81.8 | 201 |
| VGG13_BN | 6.31 | 14.8 | 68.0 | 166 |
| VGG16_BN | 4.96 | 12.0 | 58.5 | 140 |
| VGG19_BN | 51.4 | 121 |
Remarque: Torch2trt dépend de l'API Tensorrt Python. Sur Jetson, cela est inclus avec le dernier jetpack. Pour le bureau, veuillez suivre le guide d'installation de Tensorrt. Vous pouvez également essayer d'installer Torch2trt à l'intérieur de l'un des conteneurs Docker Pytorch NGC pour le bureau ou Jetson.
Pour installer la bibliothèque Torch2Trt Python, appelez ce qui suit
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py installPour installer la bibliothèque des plugins Torch2Trt, appelez ce qui suit
cmake -B build . && cmake --build build --target install && ldconfigCela comprend la prise en charge de certaines couches qui peuvent ne pas être prises en charge nativement par Tensorrt. Une fois cette bibliothèque trouvée dans le système, les convertisseurs de couche associés dans TORCH2TRT sont implicitement activés.
Remarque: Torch2trt maintient désormais les plugins en tant que bibliothèque indépendante compilée avec CMake. Cela rend les moteurs Tensorrt compilés plus portables. Si nécessaire, les plugins dépréciés (qui dépendent de Pytorch) peuvent toujours être installés en appelant
python setup.py install --plugins.
Pour installer Torch2Trt avec des fonctionnalités de la communauté expérimentale a contribué sous torch2trt.contrib , comme la formation de connaissance de quantification (QAT) ( requires TensorRT>=7.0 ), appelez ce qui suit,
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt/scripts
bash build_contrib.sh Cela vous permet d'exécuter l'exemple QAT situé ici.
Ce convertisseur fonctionne en attachant des fonctions de conversion (comme convert_ReLU ) en appels fonctionnels pytorch originaux (comme torch.nn.ReLU.forward ). Les exemples de données d'entrée sont passés par le réseau, tout comme avant, sauf que maintenant chaque fois qu'une fonction enregistrée ( torch.nn.ReLU.forward ) est rencontrée, le convertisseur correspondant ( convert_ReLU ) est également appelé par la suite. Le convertisseur est passé les arguments et l'instruction de retour de la fonction Pytorch d'origine, ainsi que le réseau Tensorrt qui est en cours de construction. Les tenseurs d'entrée de la fonction Pytorch d'origine sont modifiés pour avoir un attribut _trt , qui est le homologue Tensorrt du tenseur pytorch. La fonction de conversion utilise ce _trt pour ajouter des couches au réseau Tensorrt, puis définit l'attribut _trt pour les tenseurs de sortie pertinents. Une fois le modèle entièrement exécuté, les retours tenseurs finaux sont marqués comme des sorties du réseau Tensorrt et le moteur Tensorrt optimisé est construit.
Ici, nous montrons comment ajouter un convertisseur pour le module ReLU à l'aide de l'API Tensorrt Python.
import tensorrt as trt
from torch2trt import tensorrt_converter
@ tensorrt_converter ( 'torch.nn.ReLU.forward' )
def convert_ReLU ( ctx ):
input = ctx . method_args [ 1 ]
output = ctx . method_return
layer = ctx . network . add_activation ( input = input . _trt , type = trt . ActivationType . RELU )
output . _trt = layer . get_output ( 0 ) Le convertisseur prend un argument, un ConversionContext , qui contiendra ce qui suit
ctx.network - le réseau Tensorrt qui est en cours de construction.
ctx.method_args - Arguments de position qui ont été transmis à la fonction Pytorch spécifiée. L'attribut _trt est défini pour les tenseurs d'entrée pertinents.
ctx.method_kwargs - arguments de mots clés qui ont été transmis à la fonction Pytorch spécifiée.
ctx.method_return - La valeur renvoyée par la fonction Pytorch spécifiée. Le convertisseur doit définir l'attribut _trt le cas échéant.
Veuillez consulter ce dossier pour plus d'exemples.
Jetbot - un robot de l'IA éducatif basé sur Nvidia Jetson Nano
Jetracer - Une voiture de course éducative AI utilisant nvidia jetson nano
JetCam - Une interface de caméra Python facile à utiliser pour Nvidia Jetson
Jetcard - une image de carte SD pour la programmation Web Projets AI avec Nvidia Jetson Nano
