torch2trt
v0.5.0
Quais modelos você está usando ou esperando usar, com Tensorrt? Sinta -se à vontade para participar da discussão aqui.
Torch2trt é um conversor pytorch para tensorrt que utiliza a API de python tensorrt. O conversor é
Fácil de usar - Converter módulos com uma única função chamada torch2trt
Fácil de estender - Escreva seu próprio conversor de camada em Python e registre -o com @tensorrt_converter
Se você encontrar um problema, informe -nos!
Observe que este conversor possui cobertura limitada de Tensorrt / Pytorch. Nós o criamos principalmente para otimizar facilmente os modelos usados no projeto JetBot. Se você achar útil o conversor com outros modelos, informe -nos.
Abaixo estão alguns exemplos de uso, para mais, confira os notebooks.
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 executar o TRTModule retornado, assim como o 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 salvar o modelo como um state_dict .
torch . save ( model_trt . state_dict (), 'alexnet_trt.pth' ) Podemos carregar o modelo salvo em um TRTModule
from torch2trt import TRTModule
model_trt = TRTModule ()
model_trt . load_state_dict ( torch . load ( 'alexnet_trt.pth' ))Testamos o conversor contra esses modelos usando o script test.sh. Você pode gerar os resultados ligando
./test.sh TEST_OUTPUT.mdOs resultados abaixo mostram a taxa de transferência no FPS. Você pode encontrar a saída bruta, que inclui latência, na pasta Benchmarks.
| 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 da API Tensorrt Python. Em Jetson, isso está incluído no mais recente JetPack. Para o desktop, siga o guia de instalação do Tensorrt. Você também pode tentar instalar o Torch2Trt dentro de um dos contêineres NGC Pytorch Docker para desktop ou jetson.
Para instalar a biblioteca Torch2trt Python, chame o seguinte
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py installPara instalar a biblioteca de plugins Torch2trt, ligue para o seguinte
cmake -B build . && cmake --build build --target install && ldconfigIsso inclui suporte para algumas camadas que podem não ser suportadas nativamente pela Tensorrt. Depois que essa biblioteca é encontrada no sistema, os conversores da camada associada no Torch2Trt são implicitamente ativados.
NOTA: O Torch2trt agora mantém os plugins como uma biblioteca independente compilada com o CMake. Isso torna os motores de tensorrt compilados mais portáteis. Se necessário, os plugins depreciados (que dependem do pytorch) ainda podem ser instalados chamando
python setup.py install --plugins.
Para instalar o Torch2trt com os recursos experimentais da comunidade contribuída em torch2trt.contrib , como o Treinamento da Quantização Ciente do Ciente requires TensorRT>=7.0
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt/scripts
bash build_contrib.sh Isso permite que você execute o exemplo do Qat localizado aqui.
Este conversor funciona anexando funções de conversão (como convert_ReLU ) às chamadas funcionais originais do Pytorch (como torch.nn.ReLU.forward ). Os dados de entrada da amostra são passados pela rede, assim como antes, exceto agora, sempre que uma função registrada ( torch.nn.ReLU.forward ) é encontrada, o conversor correspondente ( convert_ReLU ) também é chamado posteriormente. O conversor é aprovado na declaração de argumentos e retorno da função Pytorch original, bem como a rede Tensorrt que está sendo construída. Os tensores de entrada para a função Pytorch original são modificados para ter um atributo _trt , que é a contraparte do Tensorrt ao tensor Pytorch. A função de conversão usa esse _trt para adicionar camadas à rede Tensorrt e, em seguida, define o atributo _trt para tensores de saída relevantes. Depois que o modelo é totalmente executado, os retornos finais dos tensores são marcados como saídas da rede Tensorrt e o mecanismo Tensorrt otimizado é construído.
Aqui mostramos como adicionar um conversor para o módulo ReLU usando a 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 ) O conversor assume um argumento, um ConversionContext , que conterá o seguinte
ctx.network - A rede Tensorrt que está sendo construída.
ctx.method_args - argumentos posicionais que foram passados para a função Pytorch especificada. O atributo _trt está definido para tensores de entrada relevantes.
ctx.method_kwargs - Argumentos de palavras -chave que foram passadas para a função Pytorch especificada.
ctx.method_return - o valor retornado pela função Pytorch especificada. O conversor deve definir o atributo _trt , quando relevante.
Consulte esta pasta para mais exemplos.
Jetbot - um robô educacional de IA baseado no Nvidia Jetson Nano
Jetracer - um carro de corrida educacional da IA usando Nvidia Jetson Nano
Jetcam - Uma interface de câmera Python fácil de usar para Nvidia Jetson
Jetcard - uma imagem do cartão SD para programas de programação da web projetos com NVIDIA Jetson Nano
