滴滴云 LSTM 优化之路

原创蔡金平滴滴云

导读

谈到神经网络，相信是当下比较火的一个词。它的发展不是一蹴而就，而是通过各代人的智慧，经过一次一次的优化、迭代才慢慢建立起当下的各种网络结构，从最简单的 MLP，到 CNN，RNN，再到增强网络，对抗网络。每一种网络结构的诞生，都是为了解决某一类特定场景的问题。本文中涉及的 LSTM 网络，就是 RNN 网络的一种变体。

工欲善其事，必先利其器。本文通过对比几种不同的实现，逐步的建立起使用 LSTM 的方法，并针对各个实现，谈谈如何进行性能分析和优化。

特别感谢滴滴云提供的高性价比 P4 实例，让测试能够顺利进行。测试环境如下：

背景知识：什么是 LSTM

当你在阅读这篇文章的时候，你的大脑其实就是一个类似 LSTM 的执行过程了。在阅读当前文字的时候，大脑会通过之前阅读过的文字，来帮助大脑理解整个句子。也就是说，大脑对于句子、文本的理解，是有记忆过程的。而不是读一个字，忘了前面的字。

RNN 就是为了解决这类问题引入的，它通过环路，解决了信息的持久化问题，引入了记忆。

但是，针对一些特殊的场景，有时候长时间的记忆不是必须的，甚至记忆太持久，会带来一些问题。比如在翻译过程中，一句话的上下文是有限的，有时候只需要依赖最近的单词就可以。LSTM（long-short term memory），也叫做长短期记忆网络，就是为了解决这类问题。

详细解释可以阅读以下这篇文章：

https://colah.github.io/posts/2015-08-Understanding-LSTMs/

第一步：从框架中来

这里主要分析目前最主流的框架 TensorFlow 中关于 LSTM OP 的实现，先来分析看看它的性能到底如何？

网上能看到最常见的 TensorFlow LSTM OP 的代码实现，如下所示，主要调用 tf.contrib.rnn.BasicLSTMCel 实现。

先直观感受一下，性能大概在多少，上代码测试，时间在 15.94ms。这个实现时间太久了，我们接下来分析，为啥这么龟速。

先上 TensorFlow 自带的工具，Timeline 得到的结果，如何使用 Timeline 可以详见附录示例。

BasicLSTMCell Timeline 分析

可以看到一次 LSTM OP 的计算，消耗了如此多的运算，放大再来看看都是些什么操作。

可以比较清晰的看到，MatMul，BiasAdd，Sigmoid 等基础的运算，这些也正是 LSTM 这个计算单元会用到的基础操作。

上图只能看到 TensorFlow OP 层面的分析，深入到 GPU 内部会是一个什么样的状态呢。来看看 nvprof 给出的结果。nvprof 使用详见附录。

BasicLSTMCell nvprof 分析

BasicLSTMCell nvprof 分析和 Timeline 的分析基本一致，BasicLSTM 就是由很多分立的 OP 实现的，在 GPU 这边表现为，由众多的 Kernel 函数实现。放大来看看这些 Kernel 之间的执行。

从图示看到，Kernel 函数之间的间隙非常的大（Kernel 函数之间，间隙越小，计算密度越大），没能充分的利用 GPU 的计算资源。

咋一看好像没什么问题，且为了完成 Input 中 64 次的迭代，这个基础 LSTM OP，采用了循环的方式实现，看似很合理，其实很不利于充分的利用 GPU 的并行性来进行计算。仔细想想，这些独立的实现是否可以合并成一个 Kernel 进行运算。

第二步：到框架中去

通过查看 Nvidia 的手册，可以知道，其实 CUDNN 内部对于 LSTM 已经有了比较好的实现。实现的原理图如下：

从该 LSTM 计算图可以看出，Wx 的计算，其实是与时间无关的，那么为了提升计算效率，这部分其实可以合并成一个大的矩阵乘法运算。

既然，CUDNN 已经有了比较好的实现，那 TensorFlow 是否有相应的实现呢，幸运的是，开源社区已经给出了一个实现的版本 cudnn_rnn.CudnnLSTM。

老规矩，测一下时间 0.72ms，天了噜，加速了接近 30 倍。快来分析一下，到底是什么黑科技。

CudnnLSTM Timeline 分析

这个图，已经非常非常的干净了，没有那么多的 TensorFlow OP 调用，核心部分只有一个 CudnnRNN 实现。（从图中能看出来 Tensorflow 实现的一些缺陷，就是在整个 LSTM OP 的执行过程中，还是存在了一些其他 OP，而且它们之间的空隙好大，而核心的 LSTM 计算，其实才不到 0.4ms，说明还有优化空间）

TensorFlow OP 层面，我们只看到了一个 CudnnRnn 实现，再进入细节，又要采用 nvprof 进行 GPU Kernel 层面的 Profiling。

CudnnLSTM nvprof 分析