备注

您正在阅读 MMClassification 0.x 版本的文档。MMClassification 0.x 会在 2022 年末被切换为次要分支。建议您升级到 MMClassification 1.0 版本，体验更多新特性和新功能。请查阅 MMClassification 1.0 的安装教程、迁移教程以及更新日志。

ONNX 转 TensorRT（试验性的）¶

如何将模型从 ONNX 转换到 TensorRT
- 准备工作
- 使用方法
支持转换至 TensorRT 的模型列表
提示
常见问题

如何将模型从 ONNX 转换到 TensorRT¶

准备工作¶

请参照安装指南从源码安装 MMClassification。
使用我们的工具 pytorch2onnx.md 将 PyTorch 模型转换至 ONNX。

使用方法¶

python tools/deployment/onnx2tensorrt.py \
    ${MODEL} \
    --trt-file ${TRT_FILE} \
    --shape ${IMAGE_SHAPE} \
    --workspace-size {WORKSPACE_SIZE} \
    --show \
    --verify \

所有参数的说明：

model : ONNX 模型的路径。
--trt-file: TensorRT 引擎文件的输出路径。如果没有指定，默认为当前脚本执行路径下的 tmp.trt。
--shape: 模型输入的高度和宽度。如果没有指定，默认为 224 224。
--workspace-size : 构建 TensorRT 引擎所需要的 GPU 空间大小，单位为 GiB。如果没有指定，默认为 1 GiB。
--show: 是否展示模型的输出。如果没有指定，默认为 False。
--verify: 是否使用 ONNXRuntime 和 TensorRT 验证模型转换的正确性。如果没有指定，默认为False。

示例：

python tools/deployment/onnx2tensorrt.py \
    checkpoints/resnet/resnet18_b16x8_cifar10.onnx \
    --trt-file checkpoints/resnet/resnet18_b16x8_cifar10.trt \
    --shape 224 224 \
    --show \
    --verify \

支持转换至 TensorRT 的模型列表¶

下表列出了保证可转换为 TensorRT 的模型。

模型	配置文件	状态
MobileNetV2	`configs/mobilenet_v2/mobilenet-v2_8xb32_in1k.py`	Y
ResNet	`configs/resnet/resnet18_8xb16_cifar10.py`	Y
ResNeXt	`configs/resnext/resnext50-32x4d_8xb32_in1k.py`	Y
ShuffleNetV1	`configs/shufflenet_v1/shufflenet-v1-1x_16xb64_in1k.py`	Y
ShuffleNetV2	`configs/shufflenet_v2/shufflenet-v2-1x_16xb64_in1k.py`	Y

注：

以上所有模型转换测试基于 Pytorch==1.6.0 和 TensorRT-7.2.1.6.Ubuntu-16.04.x86_64-gnu.cuda-10.2.cudnn8.0 进行

提示¶

如果你在上述模型的转换中遇到问题，请在 GitHub 中创建一个 issue，我们会尽快处理。未在上表中列出的模型，由于资源限制，我们可能无法提供很多帮助，如果遇到问题，请尝试自行解决。

常见问题¶

无