如何计算 LSTM 的参数量

Alan Lee

发布于 2020-10-29 10:34:47

2.4K0

发布于 2020-10-29 10:34:47

理论上的参数量

之前翻译了 Christopher Olah 的那篇著名的 Understanding LSTM Networks，这篇文章对于整体理解 LSTM 很有帮助，但是在理解 LSTM 的参数数量这种细节方面，略有不足。本文就来补充一下，讲讲如何计算 LSTM 的参数数量。

建议阅读本文前先阅读 Understanding LSTM Networks 的原文或我的译文。

首先来回顾下 LSTM。一层 LSTM 如下：

这里的xt 实际上是一个句子的 embedding（不考虑 batch 维度），shape 一般为 [seq_length, embedding_size]。图中的A 就是 cell，xt 中的词依次进入这个 cell 中进行处理。可以看到其实只有这么一个 cell，所以每次词进去处理的时候，权重是共享的，将这个过程平铺展开，就是下面这张图了：

实际上我觉得这里 x t x_t xt 并不准确，第一个 x t x_t xt 应该指的是整句话，而第二个 x t x_t xt 应该指的是这句话中最后一个词，所以为了避免歧义，我认为可以将第一个 x t x_t xt 重命名为 x x x，第二个仍然保留，即现在 x x x 表示一句话，该句话有 t + 1 t+1 t+1 个词， x t x_t xt 表示该句话的第 t + 1 t+1 t+1 个词， t ∈ [ 0 , t ] t \in [0, t] t∈[0,t]。

一个不那么小的数被多次相乘之后会变得很小，一个不那么大的数被多次相乘之后会变得很大。所以，这也是普通 RNN 容易出现梯度消失/爆炸的问题的原因。

扯远了点。

(embedding_size + hidden_size) * hidden_size + hidden_size

一个 cell 有 4 个这样结构相同的网络，那么一个 cell 的总参数量就是直接 × 4：

((embedding_size + hidden_size) * hidden_size + hidden_size) * 4

注意这 4 个权重可不是共享的，都是独立的网络。

import tensorflow as tf

model = tf.keras.model.Sequential(
    tf.keras.layers.Embedding(1000, 128),
    tf.keras.layers.LSTM(units=64),
    tf.keras.layers.Dense(10)
)
model.summary()

输入如下：

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
embedding_1 (Embedding)      (None, None, 128)         128000    
_________________________________________________________________
lstm_1 (LSTM)                (None, 64)                49408     
_________________________________________________________________
dense_1 (Dense)              (None, 10)                650       
=================================================================
Total params: 178,058
Trainable params: 178,058
Non-trainable params: 0
_________________________________________________________________

inputs = tf.random.normal([64, 100, 128])  # [batch_size, seq_length, embedding_size]
whole_seq_output, final_memory_state, final_carry_state = tf.keras.layers.LSTM(64, return_sequences=True, return_state=True)(inputs)
print(f"{whole_seq_output.shape=}")
print(f"{final_memory_state.shape=}")
print(f"{final_carry_state.shape=}")

输出：

whole_seq_output.shape=TensorShape([32, 100, 64])  # 100 表示有 100 个词，即 100 个 time step
final_memory_state.shape=TensorShape([32, 64])
final_carry_state.shape=TensorShape([32, 64])

OK，LSTM 的参数量应该挺清晰了，欢迎在评论区留下你的想法。