人脸高清算法GFPGAN之TensorRT推理

1. 综述

最近由于做数字人项目,采用的是wav2lip + GFPGAN进行人脸面部高清,但GFPGAN模型本身比较大,所以想着使用TensorRT来代替原始的pth推理看看能否提升运行速度,于是便开始了这趟windows10之下进行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文件夹

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

进入graphsurgeon文件夹

python 复制代码
pip install graphsurgeon-0.4.6-py2.py3-none-any

进入onnx_graphsurgeon文件夹

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

进入uff文件夹

python 复制代码
pip install uff-0.6.9-py2.py3-none-any.whl

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

验证

python 复制代码
import tensorrt as trt
print(trt.__version__)

2.2 pycuda安装

进入trt_inference/package下面

python 复制代码
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 复制代码
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 复制代码
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

4.2 使用TensorRT的trt模型去跑视频生成

修改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, 生成时间为: 909.6s, 总的推理时间为: 11.8s

完整的工程代码,请看百度网盘链接: 百度网盘,提取码: gfpg

相关推荐
Dfreedom.3 分钟前
在Windows上搭建GPU版本PyTorch运行环境的详细步骤
c++·人工智能·pytorch·python·深度学习
confiself11 分钟前
AndroidWorld+mobileRL
人工智能·深度学习
aneasystone本尊20 分钟前
学习 Chat2Graph 的任务分解与执行
人工智能
嘀咕博客21 分钟前
10Web-AI网站生成器
人工智能·ai工具
西柚小萌新27 分钟前
【从零开始的大模型原理与实践教程】--第一章:NLP基础概念
人工智能·自然语言处理
嘀咕博客33 分钟前
SafeEar:浙大和清华联合推出的AI音频伪造检测框架,错误率低至2.02%
人工智能·音视频·ai工具
Hello123网站33 分钟前
FinChat-金融领域的ChatGPT
人工智能·chatgpt·金融·ai工具
嘀咕博客38 分钟前
PixVerse -免费在线AI视频生成工具
人工智能·音视频·ai工具
mit6.8241 小时前
[rStar] 解决方案节点 | `BaseNode` | `MCTSNode`
人工智能·python·算法
普蓝机器人1 小时前
AutoTrack-IR-DR200底盘仿真详解:为教育领域打造的高效机器人学习实验平台
人工智能·学习·计算机视觉·机器人·移动机器人·三维仿真导航