人脸高清算法GFPGAN之TensorRT推理

news2024/9/22 1:24:43

1. 综述

最近由于做数字人项目，采用的是wav2lip + GFPGAN进行人脸面部高清，但GFPGAN模型本身比较大，所以想着使用TensorRT来代替原始的pth推理看看能否提升运行速度，于是便开始了这趟windows1之下进行GFPGAN的trt推理的折腾之旅。

2. 环境

我会提供一个我写好GFPGAN的trt推理的完整工程包。我的环境是windows10 + cuda11.7 + cudnn 8.9.2 + TensorRT-8.5.1.7 + pycuda_cuda115 + python3.8的虚拟环境。

2.1 TensorRT的环境安装

TensorRT的环境安装参考英伟达官方TensorRT8.x下载地址

2.1.1 pip安装TensorRT文件夹中的.whl文件

进入python文件夹

conda activate py38_torch # 激活你的python3.8虚拟环境
pip install tensorrt-8.5.1.7-cp38-none-win_amd64.whl

进入graphsurgeon文件夹

pip install graphsurgeon-0.4.6-py2.py3-none-any

进入onnx_graphsurgeon文件夹

pip install onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl -i http://pypi.douban.com/simple/  --trusted-host pypi.douban.com

进入uff文件夹

pip install uff-0.6.9-py2.py3-none-any.whl

还有很重要的一步就是将TensorRT的lib所有文件复制到cuda的bin下面，如下图所示

验证

import tensorrt as trt
print(trt.__version__)

2.2 pycuda安装

进入trt_inference/package下面

pip install pycuda-2021.1+cuda115-cp38-cp38-win_amd64.whl

3. 模型转换

3.1 pth2onnx(将GFPGAN的v1.4的pth模型转换为trt)

进入trt_inference/model_transformer/onnx下面

python gfpgan2onnx.py --src_model_path GFPGANv1.4.pth  --dst_model_path gfpganv1.4.onnx --img_size 512
pip install onnx-simplifier
python -m onnxsim gfpganv1.4.onnx gfpganv1.4_sim.onnx

所以就得到了trt_inference/model_transformer/onnx/gfpganv1.4_sim.onnx

3.2 onnx2trt(将GFPGAN的onnx转化为TensorRT的trt)

进入trt_inference/model_transformer/trt下面

python gfpgan2onnx2trt.py --src_model_path ../onnx/gfpganv1.4_sim.onnx  --dst_model_path gfpganv1.4.trt

所以就得到了trt_inference/model_transformer/trt/gfpganv1.4.trt模型

4 pth和trt模型推理结果比较

进入工程文件GFPGAN-master下面

4.1 使用pytorch的pth模型去跑视频生成

修改gfpgan/utils.py下面的如图两行

然后直接在GFPGAN-master下面运行 python run.py --face_path inputs/wav2lip.mp4 --audio_path inputs/vyrxlgmx.mp3 --final_path result.mp4 --outputs_path output，则会在output下面生成result.mp4, 生成时间为: 848.8s, 总的推理时间为: 5.8s