再改RNNs，试用While循环不展开特性

RNNs的循环问题

上次改RNNs虽然保证了正确性，但是还是遗留了一个问题。这里直接截取核心的一段：

    ......
    def recurrent(self, x, h_0, w_ih, w_hh, b_ih, b_hh):
        time_step = x.shape[0]
        outputs = []
        t = 0
        h = h_0
        while t < time_step:
            h = self.cell(x[t], h, w_ih, w_hh, b_ih, b_hh)
            if self.is_lstm:
                outputs.append(h[0])
            else:
                outputs.append(h)
            t += 1
        outputs = P.Stack()(outputs)
        return outputs, h
    ......

核心问题在于RNN的循环特性，也就是每个time step的结果都要依赖上一个，一般RNN层都是以一个大算子的形式直接实现双向多层，然后Python层做个binding调用。用小算子堆叠+While循环，表现在计算图里就会是一个很深的调用栈。假设一个循环中执行的算子数为N，输入序列长度为T ，则引入一个单层单向的RNN，在While循环展开后需要N*T个算子，同样也是调用栈深度。

而MindSpore默认设置的最大调用深度为1000，一旦序列长度过长，会直接报错：

RuntimeError: mindspore/ccsrc/pipeline/jit/static_analysis/evaluator.cc:201 Eval] 
Exceed function call depth limit 1000, (function call depth: 1001, simulate call depth: 998).
It's always happened with complex construction of code or infinite recursion or loop.
Please check the code if it's has the infinite recursion or call 'context.set_context(max_call_depth=value)' to adjust this value.
If max_call_depth is set larger, the system max stack depth should be set larger too to avoid stack overflow.

此时可以适当调大max_call_depth参数，但是会引入几个问题：

1. 代码修改，哪怕只有一行
2. 编译时间过长
3. max_call_depth的设置需要多次尝试

总体来说，一般序列长度在200以下不会触发max_call_depth的限制，但是通常NLP长文本任务轻易就会到达500以上的长度，且编译时间随着文本长度的增加而线性增长，导致RNNs在长文本场景下几乎不可用。

While循环不展开

所幸1.5到1.6这段前端编译团队时间攻关了这个问题，提供了While循环不展开的写法，这里先放个链接。

使用流程控制语句 - MindSpore master documentation

昨天试用了一下，把RNNs的实现改成了While循环不展开的形态：

    def construct(self, x, h, seq_length, w_ih, w_hh, b_ih, b_hh):
        time_step = x.shape[0]
        outputs = P.Zeros()((time_step, h.shape[0], h.shape[1]), x.dtype)
 
        t = P.ScalarToTensor()(0, mstype.int64)
        while t < time_step:
            x_t = x[t]
            h = self.cell(x_t, h, w_ih, w_hh, b_ih, b_hh)
            outputs[t] = h
            t += 1
 
        if seq_length is not None:
            h = get_hidden(outputs, seq_length)
            mask = sequence_mask(seq_length, time_step)
            outputs = select_by_mask(outputs, mask)
        return outputs, h
'
运行
运行

将while循环条件变量由scalar改成Tensor，即可直接使用不展开编译，但是有以下的限制：

当 while的条件为变量时， while循环体不能展开， while循环体内的表达式都是在各个step运行时计算，因此循环体内部不能出现标量、List、Tuple等非Tensor类型的计算操作，这些类型计算操作需要在图编译时期完成，与 while在执行期进行计算的机制是矛盾的。

所以实现上不再用List保存然后Stack的方式，而是先实例化一个空的Tensor，然后逐个step的填充结果。简单测试了一下GPU环境下的编译时间和执行时间（batch size = 3, input size = 16, hidden size = 32）：

可以看到循环不展开的情况下，编译时间基本不受序列长度的影响，而执行时间上相对循环展开会有一定的损耗。但是在序列长度为500以上均需要调整max_call_depth， 当长度达到10000时，编译时长已经到了完全不可接受的程度（编译了30分钟还没有结束，手动结束进程）。

小结

作为一个NLPer，循环不展开特性是刚需中的刚需，一些有时序依赖的模型都可以用这个方式进行加速，在静态图下可以基本无视文本长度进行训练了。MindSpore想要从静态图出发逐步去支持Python语法集合来达到动静统一的目标也逐步开始有达成的痕迹，相较于Tensorflow的tf.while_loop接口设计，确实基本实现了Python语法一致。写此文之前咨询过开发的同事，对于List和Tuple等非Tensor的操作应该后续也会支持，在常用场景都覆盖的情况下，将While循环默认置为不展开，又是对易用性的一个显著提升。

最后RNNs的坑基本上算是填完了，后面会把修改后的代码上到master分支，后续有好玩/有用的特性应该还会写写，顺带预告一下下一篇，讲讲CRF和怎么用MindSpore实现。