LSTM 优化之路

本文首发于滴滴云公众号：https://mp.weixin.qq.com/s/H0AKuKn1XQDEa06jDm3Nug

本文首发于滴滴云公众号：https://mp.weixin.qq.com/s/H0AKuKn1XQDEa06jDm3Nug

导读

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

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

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

背景知识：什么是 LSTM

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

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

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

详细可以阅读以下这篇文章：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 分析

和 timeline 的分析基本一致，basic lstm 就是由很多分立的 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 分析

从上图分析，可以看到，HtoD 做完显存拷贝以后，框架居然没有及时把计算发射到 GPU 上进行计算，导致了 Latency 的急剧下降，其核心计算时间只有 100 多 us，这说明如果我们直接从 CUDNN C++ API 层面进行优化，还有很大的提升空间。

第三步：自己写框架

为了更好的控制 io 和 computer 的平衡，我们选择自己写前向框架，如下是 LSTM OP 的测试，经过测试，时间在 `0.14ms`，比 tensorflow 的实现，又快了整整 5 倍。这对于 Latency 比较苛刻的场景，至关重要。

至此，我们经过三个步骤的优化，在 LSTM OP 上面取得了不错的进展。

实现	延时(ms)	加速比
tensorflow BasicLstm	15.94	x1
tensorflow CudnnLstm	0.54	x29.5
DDF CudnnLstm	0.14	x113.8

附录

如何使用滴滴云 GPU

产品详情页：https://www.didiyun.com/production/gpu.html

滴滴云提供了比较高性价比的 p4 实例，最便宜的每小时才 3 元。滴滴云GPU机型性能评测（一）：GPU计算

进入购买链接：https://app.didiyun.com/#/dc2/add?srvType=gpu_p4

进入生产环节，这个节点结束后，即可登陆实例，尽情的玩耍了。

实验环境准备：

sudo pip install python-pip
sudo pip install --upgrade pip

install cudnn

wget http?/mirrors.intra.didiyun.com/gpu_resource/cudnn-9.0-linux-x64-v7.tgz
tar -zxvf cudnn-9.0-linux-x64-v7.tgz
sudo cp cuda/include/cudnn.h /usr/local/cuda-9.0/include/
sudo cp cuda/lib64/libcudnn* /usr/local/cuda-9.0/lib64/

install tensorflow

wget https://files.pythonhosted.org/packages/d2/e3/5600f68398b46e1a576612541be8e8fb04bdd3076aae304ab22b96b2654f/tensorflow_gpu-1.7.1-cp27-cp27mu-manylinux1_x86_64.whl
sudo pip install –ignore-installed tensorflow_gpu-1.7.1-cp27-cp27mu-manylinux1_x86_64.whl

tensorflow timeline 使用教程

嵌入 timeline 相关代码，如下

执行程序后，会生成 timeline.json 文件，将该文件导入 chrome://tracing 所在的 chrome 页面即可。

nvprof 使用教程

参考文章：

• http://www.sie.es/wp-content/uploads/2015/12/cuda-profiling-tools.pdf