蚂蚁链AIoT团队与NVIDIA合作加速AI推理

蚂蚁链 AIoT 团队与 NVIDIA 合作，将量化感知训练(QAT)技术应用于深度学习模型性能优化中，并通过 NVIDIA TensorRT 高性能推理 SDK 进行高效率部署， 通过 INT8 推理， 吞吐量提升了 3 倍， 助力蚂蚁链版权 AI 平台中的模型推理服务大幅降本增效。

依托于蚂蚁链自研的区块链和 AI 技术，以及金融级的加密算法能力和云计算能力的加持，蚂蚁链版权 AI 平台能够为数字媒体时代的内容管理方、内容平台方以及创作者提供一站式的版权保护能力。作为蚂蚁链版权平台的核心能力，AI 技术可以快速提取音视频及图像特征，在短时间内完成对相似内容的自动识别和判断。得益于 NVIDIA 完善的软硬件生态，蚂蚁链团队除了可以方便的在云端 GPU 部署深度学习模型推理服务，在进一步的合作中，双方亦成功将其 TensorRT 支持的业界领先的 INT8-QAT 技术应用到版权保护业务模型上，对 AI 模型的推理效率进行大幅度优化，实现了在几乎精度无损的情况下，单 GPU 上的吞吐量提升约 300%。

利用深度学习模型高效且自动化地进行相似内容识别是蚂蚁链版权 AI 平台的一项关键技术能力，这对 AI 模型研发提出了较高的技术挑战。一方面，多媒体版权保护业务场景中作品侵权的类型复杂而多变， 为了到达高召回率和低虚警率的效果精度，蚂蚁链团队需采用较为复杂的算法模型方案;而另一方面，版权场景中高吞吐、低成本的需求又对模型提出了性能方面的挑战。换言之，算法模型需要的存算资源必须有所限制、推理性能必须较高，而一般情况下，这与模型的效果要求是矛盾的，因为深度学习模型的效果往往与模型的尺寸和所需算力正相关。而这些复杂而又多维的业务需求，也确实给团队的算法研发和优化工作，带来了相当大的挑战。

因此，蚂蚁链团队将算法研发拆解为不同的流程以解决不同维度的业务需求。简而言之，就是优先以满足业务效果指标为目的研发模型，随后借助模型压缩技术对模型存算需求进行优化，而第二部分工作实现，则需仰赖 NVIDIA 相关软硬件生态对于高性能AI模型推理加速的优秀支持。一方面，NVIDIA GPU 所提供的强大并行算力以及 INT8 Tensor Core 提供的整型计算能力，为深度学习模型推理的高效实现奠定了基础;另一方面，NVIDIA TensorRT 8 SDK 中对模型计算图的高效融合，以及对于新型模型结构的支持和优化(例如 QAT 所采用的 QDQ 结构以及 transformer-based 模型)让我们可以在 PyTorch 模型的基础上生成高效的量化推理模型。其中，QAT 作为此项目模型性能优化的重要技术，最大的优点莫过于可以在大幅提升推理速度的同时，做到几乎完全消弭量化网络整型计算相比浮点网络数值精度差异带来的算法效果损失。而近年来，随着 NVIDIA 对相关软硬件生态的对 QAT 技术的支持日臻完善，蚂蚁链团队希望能藉由 NVIDIA 相关的软硬件生态，将 QAT 技术应用到蚂蚁链版权 AI 项目中，以满足严苛的业务需求。

基于以上挑战，蚂蚁链 AIoT 技术部与 NVIDIA 技术专家合作，引入 QAT 技术对深度学习模型进行性能优化，并将模型转换为 TensorRT 引擎部署至 NVIDIA T4 GPU 进行在线推理。

首先，团队使用 QAT 技术对算法模型以“伪量化”模式进行微调，使模型得以在 INT8 整型推理模式下在算法效果指标上几乎没有损失。具体实现方案主要包括：使用 NVIDIA PyTorch Quantization 工具在模型中一些特定位置插入量化/反量化节点，在原模型的基础上构造一个“伪量化”模型。紧接着，使用该“伪量化”模型在原任务上进行一定轮数的微调，使模型参数在受到量化误差扰动的同时依然可以收敛到一个局部最优，最终最小化量化推理对模型算法指标带来的负面影响。实践中，蚂蚁链团队通过 QAT 技术可以将 INT8 模型与单精度模型的算法指标的相对差距缩小到 0.2% 以内，做到了算法效果几乎无损的模型量化。

随后，则需要将微调完成的 QAT 模型转换为真正的 INT8 版本模型部署到线上生产环境中。这一流程主要依赖 NVIDIA 提供的 TensorRT 高性能推理框架，相较其他部署框架，TensorRT 在 NVIDIA GPU 上表现出巨大的性能优势：一方面通过 graph fusion，kernel tuning 等功能，可以自动化精简网络结构，为模型各层的不同 op 寻找最优 CUDA kernel 等优化操作;更重要的是，TensorRT 8.0 开始，添加了针对“伪量化”节点的自动化解析和融合功能，使 QAT 模型落地的门槛大大降低，让 QAT 模型真正有可能在部署环境中展现出近似 INT8 后量化模型的极高推理性能。实践中，采用 INT8-QAT 的模型做到了比单精度模型约 300% 的单位时间吞吐提升，极大提高了模型推理服务的效率;同时，量化模型更低的显存占用也为模型部署带来了更高的灵活性。

借助 NVIDIA 在高性能模型推理方面完善的软硬件生态，蚂蚁链团队得以使用 INT8-QAT 技术大幅优化蚂蚁链版权 AI 平台中模型推理服务的效率，帮助系统整体降本增效。300% 的推理速度提升，以及算法指标几乎无损的特性，让此项目的 AI 模型可以做到兼顾效果与性能，帮助蚂蚁链版权 AI 平台在业界树立技术优势。

“版权保护是一个富有挑战的技术领域，互联网内容类型多样、隐匿性强、易复制与编辑等特点对我们提出了诸多挑战。蚂蚁链从 2019 年发布鹊凿版权保护平台以来，在音视图文的侵权检索、侵权比对与定位方面做了大量深入的研究工作。我们 AIoT 团队在 2022 年世界知识产权日向公众发布了蚂蚁链版权 AI 计算引擎，可以对相似内容的提取精确到帧，直接以秒为单位反馈比对结果。这种细颗粒度的识别能力极大提高了下游工作效率，同时也需要我们充分探索和利用 NVIDIA INT8-QAT 等加速技术，在效果和性能之间达到最佳平衡。接下去我们还会面向低成本高性能的版权 AI 算法演进，让技术普惠更多的普通创作者。”蚂蚁链 AIoT 高级算法专家张晓博与钱烽表示。

审核编辑：汤梓红