python之tf.data.Dataset 与 tf.keras 的性能_编程设计

我正在使用 tf.keras(tensorflow 版本 1.9.0)进行多标签文本分类。我有一个包含 185485 个火车和 46372 个验证示例的数据集。

对于第一次尝试(在 CPU 上)，我预先填充了数据并将其提供给模型:

from tensorflow import keras 
.... 
X_train = keras.preprocessing.sequence.pad_sequences(X_train, maxlen=2000) 
 
inp = keras.Input(shape=(X_train.shape[1], )) 
x = keras.layers.Embedding(len(tk.word_index) + 1, 256, mask_zero=True)(inp) 
x = keras.layers.LSTM(128, return_sequences=False)(x) 
x = keras.layers.Dropout(0.1)(x) 
x = keras.layers.Dense(y_train.shape[1], activation="sigmoid")(x) 
model = keras.Model(inputs=inp, outputs=x) 
model.compile(loss='binary_crossentropy', 
                  optimizer='adam', 
                  metrics=['accuracy']) 
model.fit(X_train, y_train, batch_size=128, epochs=300, validation_split=0.2)

该模型每个时期需要 ~130 分钟 来训练。

然后我尝试对 Dataset 做同样的事情。由于我的数据很大，我不能适应 2GB 的限制来使用 Dataset.from_tensor_slices()，我改用 Dataset.from_generator()

X_train = keras.preprocessing.sequence.pad_sequences(X_train, maxlen=2000) 
X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=777, stratify=y_train) 
 
def gen(data, labels): 
    for x, y in zip(data, labels): 
        yield x, y 
 
train_dataset = tf.data.Dataset.from_generator( 
    lambda: gen(X_train, y_train), 
    output_types=(tf.int32, tf.int32), 
    output_shapes=([2000], [y_train.shape[1]]), 
) 
 
train_dataset = train_dataset.batch(128) 
train_dataset = train_dataset.repeat() 
 
val_dataset = ... 
 
.... 
model.fit(train_dataset, epochs=300, steps_per_epoch=len(X_train)//128, validation_data=val_dataset, 
      validation_steps=len(X_val)//128)

我预计性能大致相同，但事实并非如此，因为一个 epoch 需要 ~280 分钟 来训练。我错过了什么？如何使用数据集输入实现相同的性能？

请您参考如下方法:

在 LSTM 层尝试 unroll=True。有关详细信息，请查看 Keras 网站上的 LSTM 层文档
或在 GPU 上使用 CuDNNLSTM 进一步加速