图片分类业务推理性能探索

Source : mp.weixin.qq.com

在我们的业务中，一个基于Vision Transformer (ViT)微调的色情图片分类模型扮演着关键角色。它在线下评估中表现出色，然而，当我们将它部署到生产环境提供在线推理服务时，真正的挑战才刚刚开始。我们最初采用训练时的PyTorch框架 + Flask提供WEB服务的方案，但性能测试结果给我们敲响了警钟：单张图片的端到端平均延迟 53.64ms。令人困惑的是，GPU核心、显存等关键资源的利用率始终不高。在业务压力下，我们首先依靠增加服务副本、利用多进程并发的方法来支撑，这显然不是一个可扩展、高效率的长期解决方案。那么，如何榨干GPU的每一分性能？本文将完整复盘我们的优化之旅，记录我们如何通过七个阶段的迭代，跟随性能瓶颈的不断转移，将单张图片的平均处理时间从 53.64ms 极致压缩到 8.34ms，最终实现了超过 6.4倍的性能提升。