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

Posted 2023-02-23

【中文标题】keras 中的小型 LSTM 模型不适合我的 GPU

【英文标题】：Small LSTM model in keras does not fit my GPU

【发布时间】：2019-09-15 01:47:36

【问题描述】：

我正在 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()

I have checked 我有 15 GB 的 GPU RAM 可用，根据 my estimations 的说法，批量大小为 32 的模型应该适合 3GB 的 RAM。

但是，每当我启动训练时，服务器都会耗尽内存。

公平地说，我使用的是极长的数据序列（20000 是最大序列长度），但我希望模型在内存中象征性地展开并刚好适合。

将批量大小减少到 1 也无济于事。

发生了什么事？我怎样才能使这个模型适合内存？

编辑：我尝试将序列长度减少到 2，这确实使它适合内存。但我需要序列长度保持高。如何告诉 Tensorflow 在任何时候都不要展开网络？ （我怀疑这就是幕后发生的事情，我该如何检查是否确实如此？）

编辑：如果我删除 Softmax 层，那么内存使用会再次下降到正常范围。我认为 Softmax 层导致 Tensorflow 展开网络。虽然 TimeDistributing Softmax 并没有帮助。

【问题讨论】：

请包含model.summary()的输出，参数的数量可能很大。试验不同的序列长度，看看参数数量的变化。

@MatiasValdenegro 参数的数量在我通常使用的范围内：1,474,404 参数

【参考方案1】：

为 CuDNNLSTM 层更改 LSTM 层就成功了！

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)