torch2trt
v0.5.0
¿Qué modelos estás usando o esperando usar, con Tensorrt? Siéntase libre de unirse a la discusión aquí.
TORCH2TRT es un convertidor Pytorch a Tensorrt que utiliza la API Tensorrt Python. El convertidor es
Fácil de usar: Convertir módulos con una sola función de llamada torch2trt
Fácil de extender: escriba su propio convertidor de capa en Python y regístice con @tensorrt_converter
Si encuentra un problema, ¡háganoslo saber!
Tenga en cuenta que este convertidor tiene una cobertura limitada de Tensorrt / Pytorch. Lo creamos principalmente para optimizar fácilmente los modelos utilizados en el proyecto JetBot. Si encuentra útil el convertidor con otros modelos, háganoslo saber.
A continuación se presentan algunos ejemplos de uso, para más consulte los cuadernos.
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 ]) Podemos ejecutar el TRTModule devuelto al igual que el modelo Pytorch original
y = model ( x )
y_trt = model_trt ( x )
# check the output against PyTorch
print ( torch . max ( torch . abs ( y - y_trt ))) Podemos guardar el modelo como un state_dict .
torch . save ( model_trt . state_dict (), 'alexnet_trt.pth' ) Podemos cargar el modelo guardado en un TRTModule
from torch2trt import TRTModule
model_trt = TRTModule ()
model_trt . load_state_dict ( torch . load ( 'alexnet_trt.pth' ))Probamos el convertidor contra estos modelos utilizando el script test.sh. Puede generar los resultados llamando
./test.sh TEST_OUTPUT.mdLos resultados a continuación muestran el rendimiento en FPS. Puede encontrar la salida sin procesar, que incluye latencia, en la carpeta de puntos de referencia.
| Modelo | 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 |
Nota: Torch2Trt depende de la Tensorrt Python API. En Jetson, esto se incluye con el último Jetpack. Para el escritorio, siga la guía de instalación de Tensorrt. También puede intentar instalar Torch2Trt dentro de uno de los contenedores NGC Pytorch Docker para escritorio o jetson.
Para instalar la biblioteca Torch2Trt Python, llame a lo siguiente
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py installPara instalar la biblioteca Torch2Trt Plugins, llame a lo siguiente
cmake -B build . && cmake --build build --target install && ldconfigEsto incluye soporte para algunas capas que pueden no ser compatibles de forma nativa por Tensorrt. Una vez que esta biblioteca se encuentra en el sistema, los convertidores de capa asociados en TORCH2TRT están implícitamente habilitados.
Nota: Torch2Trt ahora mantiene complementos como una biblioteca independiente compilada con CMake. Esto hace que los motores Tensorrt compilados sean más portátiles. Si es necesario, los complementos desaprobados (que dependen de Pytorch) aún pueden instalarse llamando
python setup.py install --plugins.
Instalar TORCH2TRT con características aportadas de la comunidad experimental bajo torch2trt.contrib , como la capacitación de cuantización consciente (QAT) ( requires TensorRT>=7.0 ), llame a lo siguiente,
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt/scripts
bash build_contrib.sh Esto le permite ejecutar el ejemplo QAT ubicado aquí.
Este convertidor funciona adjuntando las funciones de conversión (como convert_ReLU ) a las llamadas funcionales originales de Pytorch (como torch.nn.ReLU.forward ). Los datos de entrada de muestra se pasan a través de la red, al igual que antes, excepto ahora cada vez que se encuentra una función registrada ( torch.nn.ReLU.forward ), el convertidor correspondiente ( convert_ReLU ) también se llama después. El convertidor se pasa por los argumentos y la declaración de retorno de la función Pytorch original, así como la red tensor que se está construyendo. Los tensores de entrada a la función Pytorch original se modifican para tener un atributo _trt , que es la contraparte de tensor al tensor Pytorch. La función de conversión utiliza este _trt para agregar capas a la red Tensorrt, y luego establece el atributo _trt para tensores de salida relevantes. Una vez que el modelo se ejecuta completamente, los rendimientos de los tensores finales se marcan como salidas de la red tensor, y se construye el motor tensor optimizado.
Aquí mostramos cómo agregar un convertidor para el módulo ReLU usando la 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 ) El convertidor toma un argumento, un ConversionContext , que contendrá lo siguiente
ctx.network : la red Tensorrt que se está construyendo.
ctx.method_args - Argumentos posicionales que se pasaron a la función Pytorch especificada. El atributo _trt se establece para tensores de entrada relevantes.
ctx.method_kwargs : argumentos de palabras clave que se pasaron a la función Pytorch especificada.
ctx.method_return : el valor devuelto por la función pytorch especificada. El convertidor debe establecer el atributo _trt donde sea relevante.
Consulte esta carpeta para obtener más ejemplos.
JETBOT - Un robot de IA educativo basado en Nvidia Jetson Nano
JETRACER - Un autoaprecar de IA educativo usando Nvidia Jetson Nano
JetCam: una interfaz de cámara Python fácil de usar para Nvidia Jetson
JetCard: una imagen de tarjeta SD para proyectos de IA de programación web con Nvidia Jetson Nano
