近年来,随着深度学习领域动态 shape 推理场景激增,如机器翻译、语音识别等,针对动态 shape 模型的优化成为了热门研究议题。本文主要介绍了 DietCode 的整体架构设计,尝试了解 DietCode 如何针对动态 shape 进行调优。以 BERT 模型为例,DietCode 与 Ansor 相比, 搜索时间减少 5.88倍,性能提高约 70%。DietCode项目已经开始合入TVM中,但目前来看,只是在 tir 优化 Pass 中增加了 LocalPad,其他 feature 还未合入。
1 背景
在机器学习领域,提高计算密集型算子的性能是一项极具挑战的工作,目前业界主要有两种解决方案:使用供应商的算子库或使用 auto-scheduler 来完成算子调优。然而,前者需要大量的工程工作,算子开发周期长,限制多且专业性强;后者在 TVM 中只支持输入为静态形状的工作负载,以便模型编译时获得形状大小,如果输入为动态形状,将会导致极长的调优时间,这是因为整个搜索空间与形状强相关,auto-scheduler 不能为所有可能形状的工作负载构造一个确定的搜索空间。为了解决以上问题,DietCode 应运而生。
目前,TVM中 auto-scheduler(如:Ansor)的整体流程,如图1(a)所示,详细信息请参考Ansor: 为深度学习生成高性能张量程序。它的的输入为 TE&#x