首页 > 解决方案 > Tensorflow 1 与 Tensorflow 2 Keras 推理速度相差 2 倍以上

问题描述

我试图找出速度与两种不同型号不同的原因。

  1. 使用 tensorflow 1.x 构建的 LSTM RNN 模型:
    self.input_placeholder = tf.placeholder(
        tf.int32, shape=[self.config.batch_size, self.config.num_steps], name='Input')
    self.labels_placeholder = tf.placeholder(
        tf.int32, shape=[self.config.batch_size, self.config.num_steps], name='Target')
    embedding = tf.get_variable(
        'Embedding', initializer = self.embedding_matrix, trainable = False)
    inputs = tf.nn.embedding_lookup(embedding, self.input_placeholder)
    inputs = [tf.squeeze(x, axis = 1) for x in tf.split(inputs, self.config.num_steps, axis = 1)]
    self.initial_state = tf.zeros([self.config.batch_size, self.config.hidden_size])
    lstm_cell = tf.contrib.rnn.BasicLSTMCell(self.config.hidden_size)
    outputs, _ = tf.contrib.rnn.static_rnn(
        lstm_cell, inputs, dtype = tf.float32, 
        sequence_length = [self.config.num_steps]*self.config.batch_size)
    with tf.variable_scope('Projection'):
        proj_U = tf.get_variable('Matrix', [self.config.hidden_size, self.config.vocab_size])
        proj_b = tf.get_variable('Bias', [self.config.vocab_size])
        outputs = [tf.matmul(o, proj_U) + proj_b for o in rnn_outputs]
  1. 使用 tensorflow 2.0 keras 构建的相同模型(至少根据我的理解):
def setup_model():
    model = Sequential()
    model.add(Embedding(input_dim=vocab_size,
                        output_dim=embedding_dim,
                        weights=[embedding_matrix],
                        input_length=4,
                        trainable=False))
    model.add(LSTM(config.hidden_size, activation="tanh"))
    model.add(Dense(vocab_size, activation="softmax"))
    return model

架构是:

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
embedding (Embedding)        (None, 4, 100)            55400     
_________________________________________________________________
lstm (LSTM)                  (None, 100)               80400     
_________________________________________________________________
dense (Dense)                (None, 554)               55954     
=================================================================
Total params: 191,754
Trainable params: 136,354
Non-trainable params: 55,400
_________________________________________________________________

我期待类似的推理运行时,但使用 tensorflow 1.x 构建的推理运行时要快得多。我试图仅使用本机 tensorlow 函数将 tensorflow 1.x 模型转换为 tensorflow 2,但由于 tensorflow 从 1.x 到 2 的巨大变化,我无法转换,我只能使用 tf.keras 创建它.

在速度方面,由于我同时使用生成文本序列 + 获取单词概率,所以我没有单一的推理时间差(我无法从 tensorflow 1.x 模型修改现有 API 来获得这个)。但总的来说,我的用例至少有 2 倍的时间差异。

这种推理速度差异背后的可能原因是什么?如果需要,我很乐意提供更多信息。

标签: performancetensorflowkeras

解决方案


推荐阅读