Tensorflow RNN中的坑

由于多个版本的积累，Tensorflow中的RNN比较杂乱，到底哪个版本实际效率高，经过实测我发现和api中说明的并不一致，在此记录一下。

注意，以下相关代码和结论均运行在tensorflow==1.10上

1.脉络梳理

在1.10版本的tensorflow中，有关rnn的部分一般在以下四个包中，

• tf.contrib.rnn
• tf.contrib.cudnn_rnn
• tf.nn.rnn_cell
• tf.compat.v1.nn.rnn_cell

其中tf.nn.rnn_cell、tf.compat.v1.nn.rnn_cell和tf.contrib.rnn互相等价，那么简化为两部分：

• tf.contrib.rnn
• tf.contrib.cudnn_rnn

在此我们只考察常用的RNN cell（例如 lstm、gru等），并把等价的做个性能对比，上面包中能找到的分类如下：

LSTM：

　　tf.contrib.rnn.LSTMCell

　　tf.contrib.rnn.LSTMBlockCell

　　tf.contrib.rnn.LSTMBlockFusedCell

　　tf.contrib.cudnn_rnn.CudnnCompatibleLSTMCell

GRU:

　　tf.contrib.rnn.GRUCell

　　tf.contrib.rnn.GRUBlockCellV2

　　tf.contrib.cudnn_rnn.CudnnCompatibleGRUCell

SRU：

　　tf.contrib.rnn.SRUCell

 

2.评测代码

代码如下，除了简单的session.run计时外，还做了利用timeline工具profiling。但在后续分析时，主要还是利用计时结果。因为timeline的分析结果里面细化到每个最小操作符，比较麻烦。

因为输入数据是使用numpy随机生成的，直接copy过去就能运行， 代码中比较了几种情况，分别是lstm,bi-lstm,2 layer bi-lstm。

# -*- coding: utf-8 -*-
import tensorflow as tf
import numpy as np
import time
import functools
import json
from tensorflow.python.client import timeline


def update_timeline(chrome_trace, _timeline_dict):
    # convert crome trace to python dict
    chrome_trace_dict = json.loads(chrome_trace)
    # for first run store full trace
    if _timeline_dict is None:
        _timeline_dict = chrome_trace_dict
    # for other - update only time consumption, not definitions
    else:
        for event in chrome_trace_dict['traceEvents']:
            # events time consumption started with 'ts' prefix
            if 'ts' in event:
                _timeline_dict['traceEvents'].append(event)
    return _timeline_dict


batch_size = 1
time_step = 70
hidden_num = 512
stack_num = 1

# 设置随机种子，保持每次随机数一致
np.random.seed(0)
# 创建正态分布输入数据，batch_size*time_step*hidden_num
np_input_data = np.random.randn(batch_size, time_step, hidden_num).astype(np.float32)
np_input_len = [time_step]*batch_size
# print np_input_data

# LSTM cells
#rnn_cell = tf.contrib.rnn.LSTMCell  # child of RNNCell
rnn_cell = tf.contrib.rnn.LSTMBlockCell  # child of RNNCell
# rnn_cell = tf.contrib.rnn.LSTMBlockFusedCell  # not child of RNNCell
# rnn_cell = tf.contrib.cudnn_rnn.CudnnCompatibleLSTMCell  # child of RNNCell

# GRU cells
# rnn_cell = tf.contrib.rnn.GRUCell
# rnn_cell = tf.contrib.rnn.GRUBlockCellV2
# rnn_cell = tf.contrib.cudnn_rnn.CudnnCompatibleGRUCell

# SRU cells
# rnn_cell = tf.contrib.rnn.SRUCell

# 构建一个简单的双层双向lstm网络
input_data = tf.placeholder(dtype=tf.float32, shape=[batch_size, time_step, hidden_num], name='input_data')
trans_data = tf.transpose(input_data, [1, 0, 2])

outputs = [trans_data]
for i in range(stack_num):
    fw_rnn = rnn_cell(hidden_num,name='fw_cell_%d' % i)
    #bw_rnn = rnn_cell(hidden_num,name='bw_cell_%d' % i)
    if rnn_cell is not tf.contrib.rnn.LSTMBlockFusedCell:
        fw_rnn = tf.contrib.rnn.FusedRNNCellAdaptor(fw_rnn, use_dynamic_rnn=False)
        #bw_rnn = tf.contrib.rnn.FusedRNNCellAdaptor(bw_rnn, use_dynamic_rnn=True)
    #bw_rnn = tf.contrib.rnn.TimeReversedFusedRNN(bw_rnn)
    outputs1, state1 = fw_rnn(outputs[-1], sequence_length=np_input_len, dtype=tf.float32)
    #outputs2, state2 = bw_rnn(outputs[-1], sequence_length=np_input_len, dtype=tf.float32)
    #next_layer_input = tf.concat([outputs1, outputs2], axis=-1)
    #outputs.append(next_layer_input)

    outputs.append(outputs1)

total_time = 0
_timeline_dict = None
runs = 1000
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
for i in range(0, runs):
    # options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    run_metadata = tf.RunMetadata()
    t1 = time.time()
    # result = sess.run([outputs[-1]], feed_dict={input_data: np_input_data}, options=options,
    #                   run_metadata=run_metadata)
    result = sess.run([outputs[-1]], feed_dict={input_data: np_input_data},
                      run_metadata=run_metadata)
    t2 = time.time()
    total_time += (t2 - t1)
    # fetched_timeline = timeline.Timeline(run_metadata.step_stats)
    # chrome_trace = fetched_timeline.generate_chrome_trace_format()
    #print t2 - t1
    # _timeline_dict = update_timeline(chrome_trace, _timeline_dict)

print rnn_cell
print 'average time %f ms' % (total_time / float(runs)*1000.0)
# with open('fused_%d_runs.json' % runs, 'w') as f:
#     json.dump(_timeline_dict, f)

3.性能分析

在V100上测试，得到的结果如下

 

 其中可以看出，无论是cudnn优化版本，还是原始版本，在gpu上的表现都不尽如人意,反倒是LSTMBlockFusedCell这个实现，效果出奇的好，甚至在GPU上好于GRU和SRU，CPU上差异也不大。至于造成这些现象的原因，需要阅读这几种实现的底层源码，因为时间精力有限，我没有仔细阅读，如果有感兴趣的读者，可以深入研究。总结如下：

1.如果使用tf1中的rnn，建议使用LSTMBlockFusedCell

2.所有RNN相关实现在GPU上性能，基本和层数、方向是线性关系

3.所有RNN相关实现在CPU上性能，和层数、方向不一定是线性关系

4.GRU和SRU的实现在CPU上比GPU上效率高

5.tf1中的RNN坑很多，要谨慎使用