torch2trt
v0.5.0
Tensorrt와 함께 어떤 모델을 사용하거나 사용하고 싶습니까? 여기에서 논의에 참여하십시오.
Torch2trt는 Tensorrt Python API를 사용하는 Pytorch to Tensorrt 변환기입니다. 변환기는입니다
사용하기 쉬우기 - 단일 기능으로 모듈을 변환 torch2trt 호출
쉽게 확장 - 파이썬에 자신의 레이어 변환기를 쓰고 @tensorrt_converter 에 등록하십시오.
문제가있는 경우 알려주십시오!
이 컨버터는 Tensorrt / 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) | 나노 (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 |
참고 : Torch2trt는 Tensorrt Python API에 따라 다릅니다. Jetson에서는 이것이 최신 Jetpack에 포함되어 있습니다. 데스크탑의 경우 Tensorrt 설치 안내서를 따르십시오. 데스크탑 또는 제트슨 용 NGC Pytorch Docker 컨테이너 중 하나에 Torch2trt를 설치할 수도 있습니다.
Torch2trt Python 라이브러리를 설치하려면 다음에 전화하십시오
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py installTorch2trt 플러그인 라이브러리를 설치하려면 다음에 전화하십시오
cmake -B build . && cmake --build build --target install && ldconfig여기에는 Tensorrt에 의해 기본적으로지지되지 않을 수있는 일부 층에 대한 지원이 포함됩니다. 이 라이브러리가 시스템에서 발견되면 Torch2trt의 관련 층 변환기가 암시 적으로 활성화됩니다.
참고 : Torch2trt는 이제 CMAKE로 컴파일 된 독립적 인 라이브러리로 플러그인을 유지 관리합니다. 이를 통해 컴파일 된 Tensorrt 엔진을보다 휴대용으로 만듭니다. 필요한 경우
python setup.py install --plugins호출하여 더 이상 사용되지 않은 플러그인 (Pytorch에 따라)을 설치할 수 있습니다.
Experimental Community와 함께 Torch2trt를 설치하려면 torch2trt.contrib 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 )도 나중에 호출되는 경우를 제외하고는 이전과 마찬가지로 네트워크를 통해 전달됩니다. 컨버터는 원래 Pytorch 함수의 인수 및 반환 명령문과 구성중인 Tensorrt 네트워크가 전달됩니다. 원래 Pytorch 함수에 대한 입력 텐서는 Pytorch 텐서의 텐서에 대응하는 속성 _trt 갖도록 수정됩니다. 변환 함수는이 _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 ) 변환기는 하나의 인수 인 ConversionContext 취하며 다음을 포함합니다.
ctx.network 구성중인 Tensorrt 네트워크.
ctx.method_args 지정된 pytorch 함수로 전달 된 위치 인수. _trt 속성은 관련 입력 텐서에 대해 설정됩니다.
ctx.method_kwargs 지정된 Pytorch 함수로 전달 된 키워드 인수.
ctx.method_return 지정된 pytorch 함수에 의해 리턴 된 값. 변환기는 관련된 경우 _trt 속성을 설정해야합니다.
더 많은 예제는이 폴더를 참조하십시오.
JETBOT- Nvidia Jetson Nano를 기반으로 한 교육 AI 로봇
Jetracer- Nvidia Jetson Nano를 사용한 교육 AI Racecar
JETCAM- NVIDIA JETSON 용 사용하기 쉬운 파이썬 카메라 인터페이스
JETCARD- NVIDIA JETSON NANO를 사용한 웹 프로그래밍 AI 프로젝트를위한 SD 카드 이미지
