一个易于使用的PyTorch到TensorRT转换器（torch2trt）

【实例简介】

torch2trt是一个PyTorch到TensorRT的转换器，利用了TensorRT Python API。它易于使用，可以通过单个函数调用torch2trt来转换模块。

用法

以下是一些用法示例，更多内容请查看笔记本。

转换

import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet

# 创建一些常规的pytorch模型...
model = alexnet(pretrained=True).eval().cuda()

# 创建示例数据
x = torch.ones((1, 3, 224, 224)).cuda()

# 转换为TensorRT并将示例数据作为输入
model_trt = torch2trt(model, [x])

【实例截图】
【核心代码】
文件清单
└── torch2trt-4e820ae31b4e35d59685935223b05b2e11d47b03
    ├── benchmarks
    │   ├── JETSON_NANO.md
    │   └── JETSON_XAVIER.md
    ├── build.py
    ├── CHANGELOG.md
    ├── CLA.pdf
    ├── CMakeLists.txt
    ├── CONTRIBUTING.md
    ├── CONTRIBUTORS.md
    ├── docker
    │   ├── 21-06
    │   │   ├── build.sh
    │   │   ├── Dockerfile
    │   │   └── run.sh
    │   ├── 21-08
    │   │   ├── build.sh
    │   │   ├── Dockerfile
    │   │   └── run.sh
    │   ├── 21-09
    │   │   ├── build.sh
    │   │   ├── Dockerfile
    │   │   └── run.sh
    │   └── l4t-35.1.0
    │       ├── build.sh
    │       ├── Dockerfile
    │       └── run.sh
    ├── docs
    │   ├── benchmarks
    │   │   ├── jetson_nano.md
    │   │   └── jetson_xavier.md
    │   ├── CHANGELOG.md
    │   ├── CONTRIBUTING.md
    │   ├── css
    │   │   └── version-select.css
    │   ├── getting_started.md
    │   ├── images
    │   │   ├── chart.svg
    │   │   └── check.svg
    │   ├── index.md
    │   ├── js
    │   │   └── version-select.js
    │   ├── see_also.md
    │   └── usage
    │       ├── basic_usage.md
    │       ├── custom_converter.md
    │       └── reduced_precision.md
    ├── examples
    │   ├── contrib
    │   │   ├── pre_py3.7
    │   │   │   └── fix-getitem.patch
    │   │   └── quantization_aware_training
    │   │       ├── datasets
    │   │       │   ├── cifar10.py
    │   │       │   └── __init__.py
    │   │       ├── infer.py
    │   │       ├── __init__.py
    │   │       ├── models
    │   │       │   ├── __init__.py
    │   │       │   ├── models.py
    │   │       │   └── resnet.py
    │   │       ├── parser.py
    │   │       ├── README.md
    │   │       ├── setup.py
    │   │       ├── train.py
    │   │       └── utils
    │   │           ├── __init__.py
    │   │           ├── pytorch_nvidia_quantization.patch
    │   │           └── utilities.py
    │   ├── easyocr
    │   │   ├── download_images.sh
    │   │   ├── generate_data.py
    │   │   ├── optimize_detector.py
    │   │   ├── optimize_recognizer.py
    │   │   ├── README.md
    │   │   └── run_end2end.py
    │   ├── image_classification
    │   │   ├── conversion.ipynb
    │   │   ├── imagenet_labels.json
    │   │   └── live_demo.ipynb
    │   └── image_segmentation
    │       └── conversion.ipynb
    ├── LICENSE.md
    ├── mkdocs.yml
    ├── plugins
    │   └── src
    │       ├── example_plugin.cu
    │       ├── example_plugin.h
    │       ├── example_plugin_test.cpp
    │       ├── reflection_pad_2d_plugin.cu
    │       ├── reflection_pad_2d_plugin.h
    │       ├── reflection_pad_2d_plugin_test.cpp
    │       └── tests.cpp
    ├── README.md
    ├── requirements
    │   ├── requirements_10.txt
    │   └── requirements_8.txt
    ├── scripts
    │   ├── build_contrib.sh
    │   ├── build_docs.sh
    │   ├── build_pre_py3.7.sh
    │   ├── dump_converters.py
    │   ├── profile_timm_models.sh
    │   ├── profile_timm.py
    │   ├── push_docs.sh
    │   ├── release_build_docs.sh
    │   ├── release_push_docs.sh
    │   ├── release_test_docs.sh
    │   └── test_docs.sh
    ├── setup.py
    ├── tests
    │   ├── converter_tests
    │   │   ├── __init__.py
    │   │   ├── test_converters.py
    │   │   └── test_getitem.py
    │   ├── feature_tests
    │   │   ├── __init__.py
    │   │   ├── test_contiguous.py
    │   │   ├── test_dataset_calibrator.py
    │   │   ├── test_dataset.py
    │   │   ├── test_dynamic_shape.py
    │   │   ├── test_flatten_dynamic.py
    │   │   ├── test_flattener.py
    │   │   ├── test_flatten_module.py
    │   │   ├── test_interpolate_dynamic.py
    │   │   ├── test_legacy_max_batch_size.py
    │   │   ├── test_save_load.py
    │   │   ├── test_tensor_ne.py
    │   │   ├── test_tensor_shape_div_batch.py
    │   │   ├── test_tensor_shape.py
    │   │   └── test_version_utils.py
    │   ├── __init__.py
    │   └── model_tests
    │       ├── __init__.py
    │       ├── timm
    │       │   ├── __init__.py
    │       │   └── test_maxvit.py
    │       └── torchvision
    │           ├── __init__.py
    │           ├── test_classification_models.py
    │           └── test_segmentation_models.py
    ├── test.sh
    └── torch2trt
        ├── contrib
        │   ├── __init__.py
        │   └── qat
        │       ├── converters
        │       │   ├── __init__.py
        │       │   ├── QuantConvBN.py
        │       │   ├── QuantConv.py
        │       │   └── QuantRelu.py
        │       ├── __init__.py
        │       ├── layers
        │       │   ├── __init__.py
        │       │   ├── quant_activation.py
        │       │   ├── quant_conv.py
        │       │   ├── README.md
        │       │   └── _utils.py
        │       └── README.md
        ├── converters
        │   ├── __init__.py
        │   ├── native_converters.py
        │   ├── plugin_converters.py
        │   └── unimplemented_converters.py
        ├── dataset_calibrator.py
        ├── dataset.py
        ├── flattener.py
        ├── flatten_module.py
        ├── __init__.py
        ├── misc_utils.py
        ├── plugins
        │   ├── group_norm.cpp
        │   ├── interpolate.cpp
        │   └── plugins.cpp
        ├── test.py
        ├── torch2trt.py
        ├── trt_module.py
        ├── utils.py
        └── version_utils.py

40 directories, 139 files