深度学习框架中的动态Shape问题

动态shape存在的挑战

前几年一直做算法，这几年换了方向，一直做AI编译器和推理框架相关的工作，遇到过很多问题，也解决了很多问题，印象最深的就是动态shape的问题，这个问题一直困恼着公司推理框架的实际落地，因为随着AI技术的发展，特别是NLP领域的发展，越来越多的模型需要dynamic shape的支持，但是动态shape会给推理框架带来很多问题。目前主流的AI编译器和推理框架其实对动态shape的支持多少都有些问题，据我了解，目前ONNXRuntime对动态shape的支持是最好的，其他的比如TVM,TensorRT,TNN等对动态的支持或多或少都有问题，为什么动态shape这么困难呢？动态shape带来的主要挑战包括下面几点：

1. IR需要支持动态模型
2. 编译优化阶段需要支持动态模型
3. 代码生成需要支持动态执行

下面我们详细分析一下上面的3点

IR需要支持动态模型

做的比较早的AI编译器和推理框架，比如TVM和XLA，他们的IR设计主要针对的是静态模型，因为静态模型的优势很明显，IR设计简单，而且很容易做优化，比如内存复用优化和算子融合，而且早期的深度学习模型主要用于计算机视觉领域，基本都是CNN，大部分模型都是静态模型，即使是目标检测之类的模型，比如YOLO和SSD需要动态的支持，其实也可以通过padding之类的方法来解决，所以从这个角度来看，静态IR的问题并不算严重，但是近几年随着NLP的发展，随着RNN,Transformer等架构的发展，对动态shape的需求越来越强烈，TVM对relay进行重构的头号目标就是要解决动态shape的问题，最近在关注TVM对relay的重构情况，没有亲自测试过现在TVM对动态的支持情况，所以TVM到底有没有彻底解决动态的问题目前还不清楚。XLA作为TensorFlow中一个很重要的编译器，也是不支持动态shape的，因为XLA的IR是一种静态IR，比如XLA中将slice算子的start和limit属性设置成了常量，无法在运行期修改，所以无法支持动态shape模型。早期的ONNX也是不支持动态shape的，跟XLA一样，将slice算子的start和limit属性设置成了常量，后来对这部分算子进行了修改后才支持动态shape。具体修改方法就是将属性修改为输入参数，这样在运行期就可以动态修改start和limit的值了。

IR的设计还存在一个问题就是目前的AI模型支持广播机制，对于广播机制来说，由于动态shape无法确定明确的shape，所以在编译期无法知道是否需要广播，也不知道广播后的shape。

编译优化阶段需要支持动态模型

由于动态shape模型在编译期无法确定准确的shape信息，所以很多优化无法实现，比如内存复用优化，算子融合

代码生成需要支持动态执行

在静态模型中，AI编译器代码生成出来的实现通常只能适用于某一个特定的shape,当shape发生改变后需要重新编译，如果shape变化范围特别大，那么系统开销会很大

上面三点就是目前AI编译器和推理框架面临的主要问题，那目前有哪些解决方案呢？

目前主要有下面几个方案：

1. Nimble，是基于TVM的一个动态shape解决方案，relax应该就是基于Nimble的，对动态shape的支持比较完善
2. DISC,阿里提出来的一个动态shape的AI编译器

下面我们看一下Nimble是如何解决的

Nimble

nimble的主要解决方案如下：

1. 提出了动态类型系统
2. 定义了一个shape function在运行期计算shape
3. 符号式代码生成

动态类型系统

Nimble中引入了any类型来解决动态shape中维度的表达。

Tensor[(1, 10, Any), float32]   // 支持： 固定rank， 1维动态
Tensor[(1, Any, Any), float32]  // 支持： 固定rank， 多维动态
Tensor[(*), float32]    // 不支持： 动态rank

在最新的relax中，并没有采用any这个类型表达动态shape，而是采用了符号类型系统，符号类型系统的优势是可以执行推理，可以执行很多any类型不支持的优化。比如在符号类型系统下就可以推断出R.Tensor[(n, 4), “f32”]需要的存储空间是R.Tensor[(n, 2), “f32”]的两倍，符号类型系统下的广播规则的推断：

(n, m) + (m) => (n, m)
(n, 1, m) + (2, m) => (n, 2, m)

由此可以看出，符号类型系统是一种比any更好的表达形式，也更容易做优化。

shape function

在运行时， shape function根据算子输入计算得到输出tensor的shape。 根据算子的不同特性， 将shape function分成3类：

1. 数据无依赖型， 输出shape只与输入Tensor的shape相关
2. 数据依赖型， 输出shape需要输入Tensor的value计算得到
3. 上边界型
   对于第1种类型的算子比较好处理，也比较容易做更多的优化，但是第2种类型的算子，比较困难，关于这部分的处理，relax的设计文档给出了详细的说明：https://github.com/tlc-pack/relax/wiki/Relax-Shape-Computation-Design

符号式代码生成

论文提出了在动态shape下算子生成的两个挑战：

1. 动态shape算子如何能达到静态算子性能水平？
2. 如果扩展现有Tuning算法来支持动态shape

针对第一个挑战， 提出了一种按照tiling因子进行kernel拆分的方式，针对第二个挑战， 提出了一套基于自定义模板的搜索算法。

其实在动态shape场景中，完全不需要代码生成，目前我们自己的解决方案就是不采用代码生成，而是采用手写算子的形式，这样就可以很好的支持动态运行了。

我们的方案

我们的动态shape解决方案主要包含下面几个方面：

1. 扩展IR为动态IR，因为之前的IR是一种静态IR，需要扩展原来的IR支持动态shape模型
2. 重构编译优化，让以前针对静态的pass也能用于动态，无法支持动态shape的pass就删除掉
3. 将代码生成修改为手写算子

目前基本可以兼容主流的CV和NLP领域的AI模型的动态推理，下面是我们的测试结果：

支持的模型	支持的动态模式
ResNet50	支持N，H，W维度动态
InceptionV3	支持N，H，W维度动态
MobileNetV2	支持N，H，W维度动态
DenseNet	支持N，H，W维度动态
MTCNN	支持N，H，W维度动态
SSD-VGG16	支持N，H，W维度动态
RetinaNet	支持N，H，W维度动态
RetinaFace	支持N，H，W维度动态
YOLOV3	支持N，H，W维度动态
YOLOV5	支持N，H，W维度动态
DBNet	支持N，H，W维度动态
FCN	支持N，H，W维度动态
UNet	支持N，H，W维度动态
CRNN-LSTM	支持H，W维度动态
SVTR	支持N，W维度动态
Transformer	支持序列长度动态
GPT2	支持序列长度动态

但是我们的动态shape解决方案并不完美，大规模测试的结果来看兼容性可以达到ONNXRuntime的90%~95%，由于ONNXRuntime并不是采用编译器的技术方案，所以对动态的支持要更好一些，总体来看，目前我们的方案的兼容性做的还不错，但是依旧有些问题需要继续解决，性能方面也需要持续优化。

针对动态shape的问题就先说这么多吧，其中有很多细节没有展开说，以后有空再展开讨论，有什么问题欢迎留言讨论。