keras中的小型LSTM模型不適合我的GPU

Question

我正在用Google Collab編程一個相對較小的LSTM模型。

作為參考，我使用TensorFlow 1.13來構建模型，並使用tensorflow.keras作為keras API。

seq_len = 20000; n_classes = 4
inputs = ll.Input(shape=(seq_len,))
x = ll.Embedding(len(word_index), 1000)(inputs)
x = ll.LSTM(units=100, activation='relu', return_sequences=True)(x)
outputs = ll.Dense(units = n_classes, activation='softmax')(x)
model = Model(inputs, outputs)
model.summary()

我檢查了是否有15 GB的GPU RAM可用，根據我的估計 ，批處理大小為32的模型應該適合3GB的RAM。

但是，每當我啟動培訓時，服務器內存就會用完。

公平地說，我正在使用非常長的數據序列（最大序列長度為20000），但是我希望模型在內存中象征性地展開並適合。

將批次大小減小為1也不會有幫助。

到底是怎么回事？ 如何使該模型適合內存？

編輯：我試圖減少序列長度為2，這確實使其適合內存。 但是我需要序列長度保持較高。 我怎樣才能告訴Tensorflow在任何時候都不展開網絡？ （我懷疑這是幕后情況，如何檢查是否確實如此？）

編輯：如果我刪除Softmax層，然后內存使用率再次下降到正常范圍。 我認為Softmax層導致Tensorflow展開網絡。 但是，時間分配Softmax並沒有幫助。

Answer 1

將LSTM層更改為CuDNNLSTM層就成功了！

inputs = ll.Input(shape=(seq_len,))
x = ll.Embedding(len(word_index), 1024)(inputs)
x = ll.CuDNNLSTM(units=100, return_sequences=True)(x)
x = ll.Dense(units = n_classes, activation='softmax')(x)
outputs = x
model = Model(inputs, outputs)