Tensorflow 2.0 Keras 的训练速度比 2.0 Estimator 慢 4 倍
Posted
技术标签:
【中文标题】Tensorflow 2.0 Keras 的训练速度比 2.0 Estimator 慢 4 倍【英文标题】:Tensorflow 2.0 Keras is training 4x slower than 2.0 Estimator 【发布时间】:2019-08-05 21:06:30 【问题描述】:我们最近为 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)
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() 中打印一行数据集时,在估计器中,它会打印出非急切的张量定义。在 Keras 中,它会打印出急切的值。通过 Keras 后端代码,当它接收到一个 tf.data.dataset 作为输入时,它会急切地(而且只是急切地)处理它,这就是我在 train_input_fn() 上使用 tf.function 时它崩溃的原因。基本上,我的猜测是 DNNClassifier 的训练速度比 Keras 快,因为它在图形模式下运行更多的数据集代码。将发布任何更新/发现。
【问题讨论】:
【参考方案1】:我认为它比较慢,因为它没有在图表上执行。为了在 TF2 中的图表上执行,您需要一个用 tf.function 装饰器装饰的函数。查看this section,了解如何重构代码。
【讨论】:
感谢您的回复。 ^^ 我同意这将是放缓的最可能原因。我只是不确定,因为我假设 Google / TF 团队已经编写了 Keras 的 fit() 和 Estimator 的 train() 以在后端实现中自动使用 tf.function 或图形代码。在一天的大部分时间里,我一直在尝试以图形模式运行代码,使用 tf.function 或 tf.compat.v1.disable_eager_execution()。在图形模式下,估计器可以正确训练,但 Keras 在 DenseFeatures() 层会遇到大量奇怪的错误。我会继续挖掘并发布更新。谢谢你。 ^^ 一直在挖掘更多,发现:“Keras 层/模型继承自 tf.train.Checkpointable 并与 @tf.function 集成,这使得可以直接检查点或从中导出 SavedModels Keras 对象。” (Effective TF2) 所以我假设代码已经在使用图形模式。我唯一的线索是我们的 feature_columns 基本上是一个 100 万词汇量的 indicator_column(categorical_column(...))。我在想也许 Keras 在这种尺寸和类型上的速度要慢得多。将继续挖掘,并将发布任何发现。【参考方案2】:对于那些(像我一样)发现这个问题并使用 Keras 的嵌入层的人:
即使有 GPU,但启用了 Eager Execution,嵌入层也始终放置在 CPU 上,从而导致速度大幅下降。
请参阅https://github.com/tensorflow/tensorflow/issues/44194,其中还包含一种解决方法。
【讨论】:
以上是关于Tensorflow 2.0 Keras 的训练速度比 2.0 Estimator 慢 4 倍的主要内容,如果未能解决你的问题,请参考以下文章
使用 TensorFlow 2.0 Alpha 时无法在 Tensorboard 中看到 keras 模型图
标准化 Keras:TensorFlow 2.0 高级 API 指南
翻译: Keras 标准化:TensorFlow 2.0 中高级 API 指南
如何在 Tensorflow 2.0 + Keras 中进行并行 GPU 推理?