【源码】基于Python的CvT模型TensorRT部署优化

项目简介

本项目由摆一摆队在TRT2022复赛中贡献，选取公开的视觉Transformer模型CvT，将其部署到TensorRT上进行优化运行。项目实现了模型从PyTorch到ONNX的转换，对ONNX模型进行必要修改，最终用TensorRT对模型优化。

项目的主要特性和功能

• 模型转换与修改：完成从PyTorch到ONNX的模型转换，针对TensorRT需求修改ONNX模型，如修改Squeeze操作指定压缩维度。
• TensorRT优化：用TensorRT对修改后的ONNX模型优化，实现推理加速并保持精度。
• 性能评估：在ImageNet数据集上评估模型精度和性能，对比PyTorch、ONNX Runtime和TensorRT的推理速度。

安装使用步骤

环境配置

• 下载Nvidia官方提供的docker镜像，创建并启动容器。
• 在容器内安装1.7.1版本的PyTorch和TorchVision，以及本项目所需依赖包。

数据准备

准备ImageNet数据集，按特定格式组织数据。

模型转换

• 加载预训练模型权重，将模型从PyTorch格式转为ONNX格式。
• 修改ONNX模型的计算图，确保TensorRT能正确解析。

TensorRT部署

使用TensorRT对修改后的ONNX模型进行解析和优化，得到优化后的模型。

性能评估

在ImageNet数据集上评估模型的精度和性能，对比PyTorch、ONNX Runtime和TensorRT的推理速度。

注意：项目代码需在已配置好的Docker容器内运行，以确保环境一致性和资源隔离。

下载地址

点击下载 【提取码: 4003】【解压密码: www.makuang.net】