【SOP】使用MMDeploy将MMAction2的模型转换为TensorRT

配置环境

• MMAction2 1.2.0 配置MMAction2环境
• CUDA 11.8 安装CUDA，cuDNN，NCCL
• cuDNN 8.4.1 安装CUDA，cuDNN，NCCL
• TensorRT GA 8.6 配置TensorRT环境
• MMDeploy pip install mmdeploy
• pytorch 2.0.1, torchvision 0.15.2 conda install -y pytorch==2.0.1 torchvision==0.15.2 pytorch-cuda=11.8 -c pytorch -c nvidia
• numpy 1.26.3
• pycuda 配置TensorRT环境（只需要推理的时候使用，仅仅转换的话不需要）

执行一个MMAction2的Demo

执行 demo/demo.py ，并使用 tsm 配置文件。

python demo/demo.py "configs/recognition/tsm/tsm_imagenet-pretrained-r50_8xb16-1x1x16-50e_kinetics400-rgb.py" "../mmaction2 Ckpt/tsm_imagenet-pretrained-r50_8xb16-1x1x16-50e_kinetics400-rgb_20220831-042b1748.pth" "demo/demo.mp4" "tools/data/kinetics/label_map_k400.txt"

将Demo转为TensorRT

Clone the mmdeploy repo: git clone -b main https://github.com/open-mmlab/mmdeploy.git

假设当前路径是：

其中 Ckpt 存放了 .pth 文件，将要生成的 .onnx 和 .engine 也会放在 Ckpt 中。.engine 就是转换好的 TensorRT 文件。

执行下面的指令完成TensorRT转换：

cd mmdeploy
python tools/deploy.py "configs/mmaction/video-recognition/video-recognition_2d_tensorrt_static-224x224.py" "../mmaction2/configs/recognition/tsm/tsm_imagenet-pretrained-r50_8xb16-1x1x16-50e_kinetics400-rgb.py" "../Ckpt/tsm_imagenet-pretrained-r50_8xb16-1x1x16-50e_kinetics400-rgb_20220831-042b1748.pth" "tests/data/arm_wrestling.mp4" --work-dir "../Ckpt/tsm_trt" --device cuda:0 --show --dump-info

会遇到报错：

Error Code 4: Internal Error (input: kMIN dimensions in profile 0 are [1,250,3,224,224] but input has static dimensions [1,160,3,224,224].)

Because the shape in mmdeploy/configs/mmaction/video-recognition/video-recognition_2d_tensorrt_static-224x224.py is default as [1, 250, 3, 224, 224], while the shape of tests/data/arm_wrestling.mp4 is [1, 160, 3, 224, 224]. The shapes should be the same.

修改 [1, 250, 3, 224, 224] 为 [1, 160, 3, 224, 224] 后即可。