推理提速 3 倍：Chinese-CLIP ONNX/TensorRT 转换与加速完整指南

推理提速 3 倍Chinese-CLIP ONNX/TensorRT 转换与加速完整指南【免费下载链接】Chinese-CLIPChinese version of CLIP which achieves Chinese cross-modal retrieval and representation generation.项目地址: https://gitcode.com/GitHub_Trending/ch/Chinese-CLIPChinese-CLIP 是中文图文跨模态检索模型。这篇记录我们在 T4 上做的部署实战把 PyTorch 模型转成 ONNX 与 TensorRT 两种格式ViT-B/16 图像侧推理时延从 11.12ms 压到 3.58ms文本侧从 12.47ms 压到 1.54msMUGE zero-shot R1 只动 ±0.1。提速是真的精度几乎无损全程只用了仓库里两个现成脚本不用改一行模型代码。先定位瓶颈Chinese-CLIP 部署慢在哪线上每次检索 图像编码器一次前向 文本编码器一次前向 一次内积内积开销可以忽略瓶颈全在两个前向。PyTorch 前向带着框架层开销算子调度、图构建、张量搬运都摊在里面QPS 一上来GPU 时间大部分花在了这些零碎上。所以加速对象就是图像和文本这两个特征提取器把它们的计算图固化下来开销立刻降一个台阶。两条路线定位不同先想清楚要哪条格式定位ONNX跨平台通用中间格式CPU、NVIDIA、Apple Silicon、边缘设备都能跑TensorRTNVIDIA 专属推理引擎吃满 Tensor Core同卡延迟最低一张表核对硬件与软件版本环境装错最费时间先对着清单过一遍项目版本 / 要求说明GPUVolta 架构及以上必须支持 FP16 Tensor Core如 T4、A10显存16GB 及以上ViT-B/16 到 ViT-H/14 的转换与推理都够用CUDA11.6与 PyTorch cu116 版本对齐cuDNN8.6.0必须与 TensorRT 版本严格匹配TensorRT8.5.2.28.5.x 全系对应 cuDNN 8.6.0ONNX 工具链onnx 1.13.0 onnxruntime-gpu 1.13.1 onnxmltools 1.11.1onnxmltools 负责 FP16 转换PyTorch1.12.1cu116建议 pip 直装 wheel别混 conda cudatoolkitpip install tensorrt8.5.2.2 onnx1.13.0 onnxruntime-gpu1.13.1 onnxmltools1.11.1 pip install torch1.12.1cu116 torchvision0.13.1cu116 pip install -r requirements.txt注意 ONNX 工具链是两条路线的前置转 TensorRT 也要先过一遍 ONNX所以哪怕最终只用 TRT这一组包也得装齐。一条命令导出 Chinese-CLIP ONNX 模型pytorch_to_onnx.py 读入一个 .pt 预训练权重产出 4 个 ONNX 文件文本/图像 × FP32/FP16后续主要用 FP16 这一组。python cn_clip/deploy/pytorch_to_onnx.py \ --model-arch ViT-B-16 \ --pytorch-ckpt-path pretrained_weights/clip_cn_vit-b-16.pt \ --save-onnx-path deploy/vit-b-16 \ --convert-text --convert-vision3 个核心参数参数作用--model-arch模型规模可选 RN50 / ViT-B-16 / ViT-L-14 / ViT-L-14-336 / ViT-H-14--pytorch-ckpt-pathPyTorch ckpt 路径必须与 model-arch 对应也可指向自训 ckpt--save-onnx-path输出路径前缀生成.txt/.img.fp16/.fp32.onnx文件--convert-text/--convert-vision是两个开关--context-length默认 52。顺带一提FP16 ONNX 附带一个.extra_file文件内部存了它的路径引用转换完别挪动目录否则加载时报找不到文件。图像侧特征提取ONNX 推理五行代码下面这段做的事会话载入图像侧引擎预处理后跑一次前向拿到未归一化的图像特征。img_session onnxruntime.InferenceSession(deploy/vit-b-16.img.fp16.onnx, providers[CUDAExecutionProvider]) image preprocess(Image.open(examples/pokemon.jpeg)).unsqueeze(0) features img_session.run([unnorm_image_features], {image: image.numpy()})[0]注意输入只吃 batch1一次处理一张图批量要自己循环。文本侧特征提取ONNX 推理五行代码文本侧同理给两条中文候选词分词填充长度 52和转换时一致算出文本特征。txt_session onnxruntime.InferenceSession(deploy/vit-b-16.txt.fp16.onnx, providers[CUDAExecutionProvider]) text clip.tokenize([杰尼龟, 妙蛙种子], context_length52) features txt_session.run([unnorm_text_features], {text: text.numpy()})[0]拿到归一化特征后和图像特征内积再 softmax 就是相似度操作和 PyTorch 版完全一样。ONNX 转 TensorRT 引擎--fp16 参数速查onnx_to_tensorrt.py 吃上一步的两个 ONNX 文件构建出两个.trt引擎python cn_clip/deploy/onnx_to_tensorrt.py \ --model-arch ViT-B-16 \ --text-onnx-path deploy/vit-b-16.txt.fp16.onnx \ --vision-onnx-path deploy/vit-b-16.img.fp16.onnx \ --save-tensorrt-path deploy/vit-b-16 \ --convert-text --convert-vision \ --fp16--fp16表示构建 FP16 精度引擎前提是卡上有 FP16 Tensor Core不开 FP16 就用--fp32构建 FP32 引擎个别对 FP16 溢出敏感的算子可以经--fp16-banned-ops指定保持 FP32。转换耗时按规模不同几分钟到十几分钟。不想自己转的话官方预训练 TensorRT 引擎基于 TensorRT 8.5.2.2 构建可直接获取模型规模图像侧引擎文本侧引擎RN50rn50.img.fp16.trtrn50.txt.fp16.trtViT-B/16vit-b-16.img.fp16.trtvit-b-16.txt.fp16.trtViT-L/14vit-l-14.img.fp16.trtvit-l-14.txt.fp16.trtViT-L/14336pxvit-l-14-336.img.fp16.trtvit-l-14-336.txt.fp16.trtViT-H/14vit-h-14.img.fp16.trtvit-h-14.txt.fp16.trt下载后放进 deploy/ 目录即可直接使用。另外提醒一点转换和运行如果不在同一台机器上运行环境的 TensorRT 库版本要和转换时保持一致否则加载引擎直接报错。TensorRTModel 最小推理示例tensorrt_utils.py 封装的 TensorRTModel 用法引擎路径传进去dict 传输入就能跑前向。from cn_clip.deploy.tensorrt_utils import TensorRTModel trt_model TensorRTModel(deploy/vit-b-16.img.fp16.trt) image preprocess(Image.open(examples/pokemon.jpeg)).unsqueeze(0).cuda() features trt_model(inputs{image: image})[unnorm_image_features]图像侧到此结束文本侧输入键是text序列长度 52规则和 ONNX 版一致。T4 实测Chinese-CLIP 推理时延与 MUGE 精度对照测试环境单卡 T416GB 显存、Xeon Platinum 8163 2.5GHz、64GB 内存全部 FP16、batch1。时延是 100 次特征提取的均值来自 speed_benchmark.py时延和 zero-shot 精度放同一张表模型规模任务PyTorchONNXTensorRTViT-B/16图像推理时延 (ms)11.124.923.58ViT-B/16文本推理时延 (ms)12.473.421.54ViT-H/14图像推理时延 (ms)35.1034.0026.98ViT-B/16MUGE zero-shot R1 (%)52.152.052.0ViT-H/14MUGE zero-shot R1 (%)63.062.962.9读法两句模型越小提速越大ViT-B/16 图像侧 3.1 倍、文本侧 8.1 倍ViT-H/14 这类大模型 TensorRT 优化空间小约 1.3 倍。R1 波动全部在 ±0.1 以内工程上可以直接忽略。部署踩坑速查三个最常见的⚠️cuDNN 与 TensorRT 版本不匹配TensorRT 8.5.x 必须配 cuDNN 8.6.0小版本错一位就直接加载报错。torch 建议走 pip 的 cu116 wheel 安装别在同一环境里混 conda 的 cudatoolkit这是 cuDNN 版本漂移的头号来源。⚠️onnxruntime 缺 CUDAExecutionProviderInferenceSession 建好了却回退 CPU多半是装成了 onnxruntime 而不是 onnxruntime-gpu装错一个包GPU 加速收益全丢。⚠️FP16 精度取舍FP16 引擎是提速主力3.58ms 对 11.12ms但动态范围比 FP32 窄个别算子会饱和漂移。转换脚本默认会把 LayerNorm 一类的 POWReduce 模式保留在 FP32一般直接用--fp16即可自训 ckpt 若发现精度掉把敏感算子加进--fp16-banned-ops再转一次。一句话选型ONNX 还是 TensorRT要跨平台部署CPU、其他 GPU 厂商、边缘设备选ONNX通用中间格式runtime 生态成熟。单张 NVIDIA 卡、追极致时延与 QPS选TensorRTFP16 引擎同卡延迟最低。【免费下载链接】Chinese-CLIPChinese version of CLIP which achieves Chinese cross-modal retrieval and representation generation.项目地址: https://gitcode.com/GitHub_Trending/ch/Chinese-CLIP创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考