torch2trt
v0.5.0
Какие модели вы используете или надеетесь использовать с помощью Tensorrt? Не стесняйтесь присоединиться к обсуждению здесь.
Torch2trt - это Pytorch в Tensorrt Converter, который использует API Python Tensorrt Python. Конвертер есть
Простой в использовании - преобразовать модули с одним функцией вызов torch2trt
Легко расширить - напишите свой собственный преобразователь слоя в Python и зарегистрируйте его на @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 ]) Мы можем выполнить возвращенный TRTModule , как и оригинальная модель Pytorch
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) | Ксавье (Тенсорт) |
|---|---|---|---|---|
| Алекснет | 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 |
| Денсенет121 | 11,5 | 41,9 | 23.0 | 168 |
| Денсенет169 | 8.25 | 33,2 | 16.3 | 118 |
| Densenet201 | 6.84 | 25.4 | 13.3 | 90.9 |
| Денсенет161 | 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 зависит от API Tensorrt Python. На Jetson это включено с последним JetPack. Для настольного компьютера, пожалуйста, следуйте руководству по установке Tensorrt. Вы также можете попробовать установить Torch2trt в одном из контейнеров NGC Pytorch Docker для рабочего стола или Jetson.
Чтобы установить библиотеку Python Torch2trt, вызовите следующее
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py installЧтобы установить библиотеку плагинов Torch2trt, вызовите следующее
cmake -B build . && cmake --build build --target install && ldconfigЭто включает в себя поддержку некоторых слоев, которые не могут быть поддержаны Tensorrt. Как только эта библиотека будет найдена в системе, соответствующие преобразователи слоя в Torch2trt неявно включены.
Примечание: Torch2trt теперь поддерживает плагины в качестве независимой библиотеки, скомпилированной с Cmake. Это делает скомпилированные тенсорт -двигатели более портативными. При необходимости, устаревшие плагины (которые зависят от Pytorch), все еще могут быть установлены путем вызова
python setup.py install --plugins.
Чтобы установить TORCH2TRT с экспериментальным сообществом, внесли функции в соответствии с torch2trt.contrib , например, обучение по знанию квантования (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 ). Преобразователь передается аргументам и возвращающимся оператором исходной функции Pytorch, а также строительную сеть Tensorrt. Входные тензоры для исходной функции Pytorch модифицированы, чтобы иметь атрибут _trt , который является аналогом Tensorrt к тензору Pytorch. Функция преобразования использует этот _trt для добавления слоев в сеть Tensorrt, а затем устанавливает атрибут _trt для соответствующих выходных тензоров. После того, как модель полностью выполнена, окончательные возвраты тензоров отмечены как выходы сети Tensorrt, и создан оптимизированный двигатель TensorRt.
Здесь мы показываем, как добавить конвертер для модуля ReLU , используя API Python Tensorrt.
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
Jetracer - образовательный гоночный автомобиль с использованием Nvidia Jetson Nano
Jetcam - простой в использовании интерфейс Python Camera для Nvidia Jetson
Jetcard - изображение SD -карты для веб -программирования AI Projects с Nvidia Jetson Nano
