TensorRT-使用TensorRT部署transformer图像重建模型MST++算法-优质算法部署项目实战.zip

## MST++ TensorRT模型加速优化 :zap: ### MST++: Multi-stage Spectral-wise Transformer for Efficient Spectral Reconstruction for TensorRT Hackathon 2022 ### 1.总述 我们优化的模型为：MST++,一个RGB-to-HSI的首个基于Transformer的图像重建模型。其原始模型的相关链接如下： <div align=center> |名称|参考连接| |-|-| ||<https://arxiv.org/abs/2111.07910>| ||<https://arxiv.org/abs/2204.07908>| ||<https://github.com/caiyuanhao1998/MST/>| ||<https://github.com/caiyuanhao1998/MST-plus-plus>| || <https://zhuanlan.zhihu.com/p/501101943>| </div> 主要采用2条主线优化该网络，`TensorRT ONNXParser`和`TensorRT API`两种方式。基于对ONNXParser用Nsight进行Profiling,并最终确定了有针对性的优化方案，包括： + 网络中存在大量的`Transpose+Reshape`操作较为耗时 + 卷积，分组卷积和转置卷积 + Self-Attention 冗余结构耗时 针对于上述问题我们简化网络结构通过TensorRT API重新搭建了网络结构，实现了`L2Norm`,`Self-Attention`,`LayerNorm`,`卷积，分组卷积，反卷积`和MST++的`S_MSA`核心结构. + L2Norm结构的耗时 + LayerNorm结构耗时 + Gelu结构耗时 针对于上述结构我们开发了一些Plugin同时借鉴已有的Plugin在TensorRT API通过合并算子引入`Normalization Plugin`,`LayerNorm plugin`和`Gelu Plugin`. 整体的优化过程如下图所示：  分别对比了该模型在Pytorch(原训练框架)，onnxruntime, TensorRT ONNXParser, TensorRT API的延迟，吞吐量和加速比以及衡量精度的绝对误差和相对误差的均值，最大值和中位数（将在下面给出详细说明） 最终： + 在延迟，吞吐量，加速比，吞吐比上基于TensorRT API和Plugin的方式在FP32和FP16上整体优于其他方式； + ONNXParser FP32模式下加速比维持在1.47-2.55,在FP16模式下加速比维持在2.49-2.67; + TensorRT API FP32模式下加速比维持在1.61-1.79,在FP16模式下加速比维持在2.36-2.99; + onnxruntime对此模型没有加速效果; + 在绝对和相对误差上，TensorRT API中因大维度的ReduceSum和Plugin的误差问题精度整体略低于ONNXParser方式; + FP32模式下的绝对误差:ONNXParser可控制在 $10^{-6}$,TensorRT API控制在$10^{-5}$; + FP16模式下的绝对误差:ONNXParser可控制在$10^{-3}$,TensorRT API控制在$10^{-2}$； + 相对误差有类似的结果; + 此外我们完成了INT8量化并分析了其速度和精度。 最后我们对本项目需要的环境，项目结构和运行方式进行说明： > :black_nib: 1.项目所需环境 <details> <summary>点我查看所需环境</summary> + 主机硬件环境：`Linux version 5.4.0-110-generic (buildd@ubuntu) (gcc version 9.4.0 (Ubuntu 9.4.0-1ubuntu1~20.04.1)) `, NVIDIA显卡： `NVIDIA A10 (24G)` + 主机软件环境：`Driver Version: 510.73.08`, `CUDA 11.6/cuDNN8.4/TRT8.4`, `Docker`和`NVIDIA-Docker` + Docker镜像的使用具体可以参考：<https://github.com/NVIDIA/trt-samples-for-hackathon-cn/blob/master/hackathon/setup.md> ```shell # 拉取镜像 nvidia-docker pull registry.cn-hangzhou.aliyuncs.com/trt2022/dev # 启动容器 nvidia-docker run -it --name trt2022 -v ~/trt2022_src:/target registry.cn-hangzhou.aliyuncs.com/trt2022/dev ``` + 容器内需要按照如下安装环境 ```shell # 安装必要的python package pip3 install -r requirements.txt -i https://mirrors.aliyun.com/pypi/simple/ # 因过程中安装了opencv，需要安装libgl1 apt-get update && apt-get install libgl1 ``` >:warning: 注: > > 1.本项目的测试结果依赖于上述环境 > > 2.本项目同时用TensorRT 8.4GA版进行了测试，经测试，测试结果的一致性是相同的，TensorRT 8.4GA镜像: `registry.cn-hangzhou.aliyuncs.com/trt2022/trt-8.4-ga` </details> > :black_nib: 2.项目结构说明 <details> <summary>点我查看项目结构</summary> ```shell . ├── architecture # MST++模型相关的python代码(主要来源https://github.com/caiyuanhao1998/MST-plus-plus) │ ├── README.md │ ├── __init__.py │ ├── edsr.py │ ├── HDNet.py │ ├── hinet.py │ ├── hrnet.py │ ├── HSCNN.py │ ├── MIRNet.py │ ├── MPRNet.py │ ├── MST.py | ├── Restormer.py # 以上python代码是MST++对比的其他模型的结构比如edsr,HDNet,...,MST │ └── MST_Plus_Plus.py # MST++模型结构 | ├── data # 测试数据，来源于MST++ GitHub项目，代码提交时具体数据不提供,详见该文件夹下的readme │ ├── ARAD_1K_0001.jpg │ └── ARAD_1K_0002.jpg ├── model # 模型存放位置，pth模型，onnx模型，模型权重文件，plan文件,提供模型`下载云盘地址`,详见该文件夹下的readme | ├── docs # README的静态资源文件存放 | ├── LayerNormPlugin # LayerNorm Plugin │ ├── layer_norm.cuh │ ├── LayerNormPlugin.cu │ ├── LayerNormPlugin.h │ └── Makefile | ├── NormalizePlugin # Normalization Plugin (L2Norm) │ ├── NormalizePlugin.cu │ ├── NormalizePlugin.h │ └── Makefile | ├── torch2onnx.py # pytorch模型转onnx模型 ├── get_weights.py # 获取模型权重文件 ├── mst_config.ini # tensorRT API模型结构配置文件 ├── calibrator.py # int8量化数据加载生成量化cache ├── mst_onnxparser.py # onnxparser方式的INT8量化生成plan文件 ├── mst_trt_api.py # TensorRT API方式生成TensorRT plan **（核心文件） ├── performance_latency.py # benchmark的计算和可视化 ├── performance_accuracy.py # 精度的计算和可视化 ├── performance_int8.py # int8量化benchmark和精度的计算和可视化 ├── build.sh # 完整运行本项目的全部命令，可以执行./build.sh运行(但是会很耗时！约4个小时) | ├── requirements.txt # Python package安装列表 ├── LICENSE └── README.md ``` </details> > :black_nib: 3.运行方式说明 <details> <summary>点我查看运行方式</summary> + 1.Pytorch模型转ONNX ```shell python3 torch2onnx.py -h # 为了测试可以生成任意batch_size：[1,2,4,8,16]（为了避免维度的复杂变换和额外结点，我们尽量避免使用Dynamic shape(个人认为能不用dynamic shape就不用！) python3 torch2onnx.py --batch_size=1 --onnx_path=./model/mst_plus_plus_b1.onnx --simplifier ``` + 2.TensorRT ONNXParser模型序列化 ```shell #FP32 trtexec --onnx=./model/mst_plus_plus_b1.onnx --saveEngine=./model/mst_plus_plus_onnxparser_b1.plan --workspace=3000 --verbose #FP16 trtexec --onnx=./model/mst_plus_plus_b1.onnx --saveEngine=./model/mst_plus_plus_onnxparser_b1_fp16.plan --workspace=3000 --verbose --fp16 ``` + 3.使用Nsight Profiling ONNXParser ```shell nsys profile -o mst_onnx trtexec --loadEngine=./model/mst_plus_plus_onnxparser_b1.plan --warmUp=0 --duration=0 --iterations=50 --plugins=./LayerNormPlugin.so ``` 关于Plugin的编译可以直接进入对应的文件夹执行`make`编译得到`so`文件,将编译好的`so`文件移动到项目根目录`mv *.so ../` + 4.TensorRT API 模型序列化 ```shell # 获取模型weights用于TensorRT API搭建网络 python3 get_weights.py -h # 我们