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坑很多，要謹慎使用