torch2trt
v0.5.0
Welche Modelle verwenden Sie oder hoffen Sie mit Tensorrt? Fühlen Sie sich frei, der Diskussion hier beizutreten.
Torch2TRT ist ein Pytorch -to Tensorrt -Konverter, der die Tensorrt -Python -API verwendet. Der Konverter ist
Einfach zu bedienend - Umwandeln Module mit einem einzelnen Funktionsaufruf torch2trt
Einfach zu erweitern - schreiben Sie Ihren eigenen Ebenenkonverter in Python und registrieren Sie ihn bei @tensorrt_converter
Wenn Sie ein Problem finden, lassen Sie es uns bitte wissen!
Bitte beachten Sie, dass dieser Konverter nur eine begrenzte Abdeckung von Tensorrt / Pytorch hat. Wir haben es hauptsächlich erstellt, um die im JETBOT -Projekt verwendeten Modelle einfach zu optimieren. Wenn Sie den Konverter mit anderen Modellen hilfreich finden, teilen Sie uns dies bitte mit.
Im Folgenden finden Sie einige Verwendungsbeispiele, um weitere Informationen zu den Notizbüchern zu erhalten.
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 ]) Wir können das zurückgegebene TRTModule genau wie das ursprüngliche Pytorch -Modell ausführen
y = model ( x )
y_trt = model_trt ( x )
# check the output against PyTorch
print ( torch . max ( torch . abs ( y - y_trt ))) Wir können das Modell als state_dict speichern.
torch . save ( model_trt . state_dict (), 'alexnet_trt.pth' ) Wir können das gespeicherte Modell in eine TRTModule laden
from torch2trt import TRTModule
model_trt = TRTModule ()
model_trt . load_state_dict ( torch . load ( 'alexnet_trt.pth' ))Wir haben den Konverter gegen diese Modelle mit dem Test.SH -Skript getestet. Sie können die Ergebnisse erzeugen, indem Sie aufrufen
./test.sh TEST_OUTPUT.mdDie folgenden Ergebnisse zeigen den Durchsatz in FPS. Sie finden die Rohausgabe, die Latenz umfasst, im Benchmarks -Ordner.
| Modell | 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 |
Hinweis: Torch2TRT hängt von der Tensorrt Python API ab. Auf Jetson ist dies im neuesten Jetpack enthalten. Für Desktop folgen Sie bitte dem Tensorrt -Installationshandbuch. Sie können auch versuchen, Torch2TRT in einem der NGC Pytorch Docker -Container für Desktop oder Jetson zu installieren.
Rufen Sie Folgendes an, um die Torch2TRT -Python -Bibliothek zu installieren
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py installRufen Sie Folgendes an, um die Torch2TRT -Plugins -Bibliothek zu installieren
cmake -B build . && cmake --build build --target install && ldconfigDies beinhaltet die Unterstützung einiger Schichten, die möglicherweise nicht von Tensorrt nativ unterstützt werden. Sobald diese Bibliothek im System gefunden wurde, sind die zugehörigen Ebenenkonverter in Torch2TRT implizit aktiviert.
HINWEIS: Torch2TRT verwaltet jetzt Plugins als unabhängige Bibliothek, die mit CMAKE zusammengestellt wurde. Dies macht kompilierte Tensorrt -Motoren tragbarer. Bei Bedarf können die veralteten Plugins (die von Pytorch abhängen) weiterhin durch Aufrufen von
python setup.py install --plugins-installiert.
Um Torch2TRT mit experimentellen Community zu installieren, haben Merkmale unter torch2trt.contrib beigetragen, wie das Quantization Awesare Training (QAT) ( requires TensorRT>=7.0 ), rufen Sie die folgenden an.
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt/scripts
bash build_contrib.sh Auf diese Weise können Sie das hier lokalisierte QAT -Beispiel ausführen.
Dieser Konverter funktioniert durch Anhängen von Konvertierungsfunktionen (wie convert_ReLU ) an die ursprünglichen Pytorch -Funktionsaufrufe (wie torch.nn.ReLU.forward ). Die Beispieleingabedaten werden wie zuvor durch das Netzwerk weitergeleitet, außer jetzt, wenn eine registrierte Funktion ( torch.nn.ReLU.forward ) auftritt, wird auch danach der entsprechende Konverter ( convert_ReLU ) aufgerufen. Der Konverter wird die Argumente und die Rückgabeerklärung der ursprünglichen Pytorch -Funktion sowie des Tensorrt -Netzwerks, das konstruiert wird, übergeben. Die Eingangstensoren für die ursprüngliche Pytorch -Funktion werden so modifiziert, dass sie ein Attribut _trt haben, das das Tensorrt -Gegenstück zum Pytorch -Tensor ist. Die Konvertierungsfunktion verwendet dieses _trt , um dem Tensorrt -Netzwerk Ebenen hinzuzufügen, und legt dann das _trt -Attribut für relevante Ausgabe -Tensoren fest. Sobald das Modell vollständig ausgeführt wurde, werden die endgültigen Tensoren als Ausgänge des Tensorrt -Netzwerks markiert und die optimierte Tensorrt -Engine erstellt.
Hier zeigen wir, wie Sie mit der Tensorrt Python API einen Konverter für das ReLU -Modul hinzufügen.
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 ) Der Konverter nimmt ein Argument, einen ConversionContext , der Folgendes enthalten wird
ctx.network - Das Tensorrt -Netzwerk, das konstruiert wird.
ctx.method_args - Positionsargumente, die an die angegebene Pytorch -Funktion übergeben wurden. Das _trt -Attribut ist für relevante Eingabetensoren festgelegt.
ctx.method_kwargs - Keyword -Argumente, die an die angegebene Pytorch -Funktion übergeben wurden.
ctx.method_return - Der von der angegebene Pytorch -Funktion zurückgegebene Wert. Der Konverter muss das _trt -Attribut gegebenenfalls festlegen.
Weitere Beispiele finden Sie in diesem Ordner.
JETBOT - Ein Bildungs -KI -Roboter, der auf Nvidia Jetson Nano basiert
Jetracer - Ein Bildungs -KI -Rennwagen mit Nvidia Jetson Nano
JETCAM - Eine einfach zu verwendende Python -Kamera -Schnittstelle für Nvidia Jetson
JetCard - Ein SD -Kartenbild für Webprogrammierprojekte mit Nvidia Jetson Nano
