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

Question

我们最近在 TF 2.0 中切换到 Keras，但是当我们将其与 2.0 上的 DNNClassifier Estimator 进行比较时，我们发现使用 Keras 的速度降低了大约 4 倍。 但是我终其一生都无法弄清楚为什么会发生这种情况。 两者的其余代码是相同的，使用返回相同 tf.data.Dataset 的 input_fn()，并使用相同的 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)

凯拉斯代码：

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)

更新：为了进一步测试，我编写了一个自定义子类模型（请参阅： 专家入门），它的运行速度比 Keras 快，但比 Estimators 慢。 如果 Estimator 在 100 秒内训练，自定义模型需要大约 180 秒，Keras 大约需要大约 350 秒。 一个有趣的提示是，Estimator 使用 Adam() 运行得比 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 中，它会打印出 Eager 值。 通过 Keras 后端代码，当它接收 tf.data.dataset 作为输入时，它会急切地（并且只是急切地）处理它，这就是为什么每当我在 train_input_fn() 上使用 tf.function 时它都会崩溃。 基本上，我的猜测是 DNNClassifier 的训练速度比 Keras 快，因为它在图形模式下运行更多的数据集代码。 将发布任何更新/发现。

Answer 1

我相信它比较慢，因为它没有在图表上执行。 为了在 TF2 中的图形上执行，您需要一个用 tf.function 装饰器装饰的函数。 查看此部分以了解有关如何重构代码的想法。

Answer 2

对于那些（像我一样）发现这个问题并使用 Keras 嵌入层的人：

即使存在 GPU，但启用了 Eager Execution，Embedding 层也始终放置在 CPU 上，从而导致大量减速。

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