Tensorflow 2.0 Keras的训练速度比2.0 Estimator慢4倍

时间:2019-03-14 20:35:06

标签: python tensorflow keras tensorflow-estimator tensorflow2.0

我们最近将TF 2.0切换为Keras,但将其与2.0上的DNNClassifier Estimator进行比较时,我们发现Keras的速度慢了约4倍。但是我无法为自己的生活弄清楚为什么会这样。两者的其余代码都相同,使用一个input_fn()返回相同的tf.data.Dataset,并使用相同的feature_columns。几天来一直在努力解决这个问题。任何帮助将不胜感激。谢谢

估算器代码:

estimator = tf.estimator.DNNClassifier(
        feature_columns = feature_columns,
        hidden_units = [64,64],
        activation_fn = tf.nn.relu,
        optimizer = 'Adagrad',
        dropout = 0.4,
        n_classes = len(vocab),
        model_dir = model_dir,
        batch_norm = false)

estimator.train(input_fn=train_input_fn, steps=400)

Keras代码:

feature_layer = tf.keras.layers.DenseFeatures(feature_columns);

model = tf.keras.Sequential([
        feature_layer,
        layers.Dense(64, input_shape = (len(vocab),), activation = tf.nn.relu),
        layers.Dropout(0.4),
        layers.Dense(64, activation = tf.nn.relu),
        layers.Dropout(0.4),
        layers.Dense(len(vocab), activation = 'softmax')]);

model.compile(
        loss = 'sparse_categorical_crossentropy',
        optimizer = 'Adagrad'
        distribute = None)

model.fit(x = train_input_fn(),
          epochs = 1,
          steps_per_epoch = 400,
          shuffle = True)

更新:为进一步测试,我编写了一个自定义的子类模型(请参见:Get Started For Experts),该模型的运行速度比Keras快,但比Estimators慢。如果Estimator训练时间为100秒,则自定义模型大约需要180秒,而Keras需要大约350秒。有趣的是,Adam()的Estimator运行速度比Adagrad()慢,而Keras的运行速度似乎更快。使用Adam(),Keras花费的时间不到DNNClassifier的两倍。假设我没有弄乱自定义代码,我开始认为DNNClassifier仅具有很多后端优化/效率,使其运行速度比Keras快。

自定义代码:

class MyModel(Model):
  def __init__(self):
    super(MyModel, self).__init__()
    self.features = layers.DenseFeatures(feature_columns, trainable=False)
    self.dense = layers.Dense(64, activation = 'relu')
    self.dropout = layers.Dropout(0.4)
    self.dense2 = layers.Dense(64, activation = 'relu')
    self.dropout2 = layers.Dropout(0.4)
    self.softmax = layers.Dense(len(vocab_of_codes), activation = 'softmax')

  def call(self, x):
    x = self.features(x)
    x = self.dense(x)
    x = self.dropout(x)
    x = self.dense2(x)
    x = self.dropout2(x)
    return self.softmax(x)

model = MyModel()
loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adagrad()

@tf.function
def train_step(features, label):
  with tf.GradientTape() as tape:
    predictions = model(features)
    loss = loss_object(label, predictions)
  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))

itera = iter(train_input_fn())
for i in range(400):
  features, labels = next(itera)
  train_step(features, labels)

更新:可能似乎是数据集。当我在train_input_fn()内的估计器中打印数据集的一行时,它会打印出非急切的Tensor定义。在Keras中,它会打印出渴望的值。通过Keras后端代码,当它接收到tf.data.dataset作为输入时,它会急切地(并且仅急切地)对其进行处理,这就是为什么当我在train_input_fn()上使用tf.function时它都崩溃的原因。基本上,我的猜测是DNNClassifier的训练速度比Keras快,因为它在图形模式下运行更多的数据集代码。将发布所有更新/发现。

2 个答案:

答案 0 :(得分:2)

我认为它比较慢,因为它没有在图形上执行。为了在TF2中的图形上执行,您需要一个使用tf.function装饰器装饰的函数。请查看this section,以获取有关如何重组代码的想法。

答案 1 :(得分:1)

对于那些(像我一样)发现这个问题并使用 Keras 的嵌入层的人:

即使存在 GPU,但启用了 Eager Execution,Embedding 层也始终放置在 CPU 上,从而导致速度大幅下降。

请参阅 https://github.com/tensorflow/tensorflow/issues/44194,其中还包含一个解决方法。

相关问题