torch2trt下載torch2trt源代碼下載

torch2trt

Python

v0.5.0

下載

TORCH2TRT

您正在使用什麼模型或希望與Tensorrt一起使用？隨時在這裡加入討論。

Torch2trt是使用張力python api的張力轉換器的pytorch。轉換器是

易於使用 - 帶有單個功能調用torch2trt模塊轉換
易於擴展 - 在Python中寫下您自己的圖層轉換器，然後將其註冊為@tensorrt_converter

如果您發現問題，請告訴我們！

請注意，此轉換器的張力 / pytorch覆蓋率有限。我們主要創建它是為了輕鬆優化Jetbot項目中使用的模型。如果您發現轉換器對其他型號有用，請告訴我們。

用法

以下是一些用法示例，有關更多信息，請查看筆記本。

轉變

 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 ])

執行

我們可以像原始的pytorch模型一樣執行返回的TRTModule

 y = model ( x )
y_trt = model_trt ( x )

# check the output against PyTorch
print ( torch . max ( torch . abs ( y - y_trt )))

保存和加載

我們可以將模型保存為state_dict 。

 torch . save ( model_trt . state_dict (), 'alexnet_trt.pth' )

我們可以將保存的模型加載到TRTModule中

 from torch2trt import TRTModule

model_trt = TRTModule ()

model_trt . load_state_dict ( torch . load ( 'alexnet_trt.pth' ))

型號

我們使用test.sh腳本對這些模型測試了轉換器。您可以通過調用來生成結果

./test.sh TEST_OUTPUT.md

下面的結果顯示了FPS中的吞吐量。您可以在基准文件夾中找到包括延遲的原始輸出。

模型	納米（pytorch）	納米（張力）	Xavier（Pytorch）	Xavier（張力）
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

設定

注意：Torch2Trt取決於張力python API。在Jetson上，最新的JetPack包含在內。對於桌面，請遵循“ Tensorrt安裝指南”。您也可以嘗試在NGC Pytorch Docker容器之一中安裝Torch2Trt，以用於台式機或Jetson。

步驟1-安裝Torch2trt Python庫

要安裝Torch2trt Python庫，請致電以下

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install

步驟2（可選） - 安裝TORCH2TRT插件庫

要安裝TORCH2TRT插件庫，請致電以下

cmake -B build . && cmake --build build --target install && ldconfig

這包括對某些層的支持，而這些層可能不受Tensorrt的本質支持。一旦在系統中找到了此庫後，就隱式啟用了Torch2Trt中關聯的層轉換器。

注意：Torch2trt現在將插件保存為帶有CMAKE的獨立庫。這使得彙編的張力引擎更便宜。如果需要，可以通過調用python setup.py install --plugins來安裝不推薦的插件（取決於Pytorch）。

步驟3（可選） - 安裝實驗社區貢獻功能

要在torch2trt.contrib下使用實驗社區的TORCH2TRT，例如量化意識培訓（QAT）（ requires TensorRT>=7.0 ），請致電以下內容，

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt/scripts    
bash build_contrib.sh

這使您可以在此處運行QAT示例。

它如何工作？

該轉換器通過將轉換功能（例如convert_ReLU ）連接到原始的pytorch函數呼叫（例如torch.nn.ReLU.forward ）來起作用。就像以前一樣，示例輸入數據通過網絡傳遞，除非遇到註冊函數（ torch.nn.ReLU.forward ），否則相應的convert_ReLU （ convert_relu ）也會在之後調用。轉換器通過原始pytorch函數的參數和返回語句以及正在構建的張力網絡。原始pytorch函數的輸入張量被修改為具有屬性_trt ，它是pytorch張量的張力。轉換函數使用此_trt將圖層添加到Tensorrt網絡中，然後為相關的輸出張量設置_trt屬性。一旦模型完全執行，最終的張量返回將標記為Tensorrt網絡的輸出，並構建了優化的Tensorrt引擎。

如何添加（或覆蓋）轉換器

在這裡，我們顯示瞭如何使用Tensorrt Python API添加ReLU模塊的轉換器。

 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 )