%20(2).svg)
项目简介
本项目聚焦于利用TensorRT对LeViT模型开展优化与推理加速工作。LeViT是Facebook Research研发的基于Transformer的图像分类模型，具备较高的推理速度与精度。项目借助FP16和INT8优化、自定义Softmax插件以及TensorRT引擎构建，达成了LeViT模型的显著加速。
项目的主要特性和功能
- 模型转换：可将PyTorch格式的LeViT模型转换为TensorRT格式。
- 性能优化：通过FP16和INT8量化以及自定义Softmax插件，大幅提升模型推理速度。
- 推理测试：提供脚本用于测试TensorRT和PyTorch模型的推理速度与精度。
- 引擎构建：支持构建TensorRT引擎，并运用自定义插件进行优化。
安装使用步骤
环境准备
- 拉取TensorRT Docker镜像：
shell docker run --gpus all -it --rm nvcr.io/nvidia/tensorrt:22.05-py3 - 拉取代码仓库（此处原总结未给出命令，需补充完整代码仓库地址）：
shell - 安装Python依赖：
shell pip3 install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu113
运行步骤
- 构建并运行Softmax插件：
shell cd LeViT ./build_engine_run_plugin.sh - 构造TensorRT引擎：
shell cd LeViT python3 trt_convert_onnx.py MODEL python3 trt_build_engine.py --onnx-path ONNX_MODEL --engine-path ENGINE_PATH [--enable-fp16] [--enable-int8] - 测试模型精度：
shell python3 valid.py --data-path IMAGENET_ROOT --model MODEL_NAME --type TYPE --engine-path ENGINE_PATH - 测试模型推理速度：
shell trtexec --loadEngine=ENGINE_PATH --shapes=input:16x3x224x224 --useSpinWait - 生成calibration数据集：
shell python3 generate_calibration_data.py --data-path IMAGENET_ROOT
下载地址
点击下载 【提取码: 4003】【解压密码: www.makuang.net】