pytorch中LSTM的细节分析理解

虽然看了一些很好的blog了解了LSTM的内部机制，但对框架中的lstm输入输出和各个参数还是没有一个清晰的认识，今天打算彻底把理论和实现联系起来，再分析一下pytorch中的LSTM实现。

先说理论部分。一个非常有名的blog把原理讲得很清楚，推荐参考。总之就是这些公式：

简单来说就是，LSTM一共有三个门，输入门，遗忘门，输出门，$i,f,o$分别为三个门的程度参数，$g$是对输入的常规RNN操作。公式里可以看到LSTM的输出有两个，细胞状态$c^{\prime }$和隐状态$h^{\prime }$，$c^{\prime }$是经输入、遗忘门的产物，也就是当前cell本身的内容，经过输出门得到$h^{\prime }$，就是想输出什么内容给下一单元。

那么实际应用时，我们并不关心细胞本身的状态，而是要拿到它呈现出的状态$h^{\prime }$作为最终输出。以pytorch中的LSTM为例：

torch.nn.LSTM(*args, **kwargs)

官方API：
https://pytorch.org/docs/stable/nn.html?highlight=lstm#torch.nn.LSTM

---

• 参数
  – input_size
  – hidden_size
  – num_layers
  – bias
  – batch_first
  – dropout
  – bidirectional

---

• 输入
  – input (seq_len, batch, input_size)
  – h_0 (num_layers * num_directions, batch, hidden_size)
  – c_0 (num_layers * num_directions, batch, hidden_size)

---

• 输出
  – output (seq_len, batch, num_directions * hidden_size)
  – h_n (num_layers * num_directions, batch, hidden_size)
  – c_n (num_layers * num_directions, batch, hidden_size)

---

用起来很简单，当作黑箱时只要设置参数让它输出我们想要的shape就行了，但这些参数好像很难和前面公式里的那些联系起来，不便于理解和灵活使用。

先看一张很好的图（LSTM神经网络输入输出究竟是怎样的？ - Scofield的回答 - 知乎）：

这张图是以MLP的形式展示LSTM的传播方式（不用管左边的符号，输出和隐状态其实是一样的），方便理解hidden_size这个参数。其实hidden_size在各个函数里含义都差不多，就是参数W的第一维（或最后一维）。那么对应前面的公式，hidden_size实际就是以这个size设置所有W的对应维。

再看另一张很好的图（https://medium.com/mlreview/understanding-lstm-and-its-diagrams-37e2f46f1714）：

这张图非常便于理解参数num_layers。实际上就是个depth堆叠，每个蓝色块都是LSTM单元，只不过第一层输入是$x_t,h_{t-1}^{(0)},c_{t-1}^{(0)}$，中间层输入是$h_t^{(k-1)},h_{t-1}^{(k)},c_{t-1}^{(k)}$。

剩下的参数就比较好理解了，input_size即输入的隐层维度，比如embedding_dim。batch_first，第一维是否是batch，为什么要设置这个参数后面再说。bidirectional，是否为双向LSTM。

接下来看一下输入输出。关于input，API中提到也可以是一个packed变量序列，这个后面再讲。输入输出中$h$和$c$的shape (num_layers * num_directions, batch, hidden_size) 也是个容易困惑的点，但有了上面那张图就好说多了。绿色块$h_n,c_n$即长度为n的序列的最终输出，可以看出是所有depth输出的拼接，维度是num_layers。双向LSTM情况，相当于有两个图中的网络，只不过输入颠倒过来了，再将这两个最终隐状态拼接起来，维度num_layers*2。

最后看一下输出output。初学时看别人的代码，总是搞不清到底是取output还是用$h_n$，怎么用的都有。其实从图中可以看到，output就是最后一个layer上，序列中每个状态$h$的集合（若为双向则按位置拼接，输出维度2*hidden_size），所以$h_n$就是$output[-1,:,:]$。我们使用LSTM的目的是得到整个序列的embedding，与序列长度无关，由于LSTM具有序列信息传递性，因此一般都取$h_n$当作序列输出。但双向LSTM推广后，每个位置的隐层输出都可以作为当前词的一个融合了上下文的embedding，因此取output再以词粒度进行融合（比如attention加权求和、pooling）也是一种主流方法。

理论终于和实践联系起来了，下面来具体分析一下pytorch的LSTM实现。

pytorch的LSTM

---

1、torch.nn.LSTMCell(input_size, hidden_size, bias=True)

官方API：https://pytorch.org/docs/stable/nn.html?highlight=lstm#torch.nn.LSTMCell**

一个LSTM单元。相当于一个time step的处理，应该是对应TensorFlow里类似的实现。基本不用。

2、torch.nn.LSTM(*args, **kwargs)

官方API：https://pytorch.org/docs/stable/nn.html?highlight=lstm#torch.nn.LSTM

前面基本讲得差不多了，只剩下两处：参数batch_first和input的packed variable length sequence。

为什么要有batch_first这个参数呢？常规的输入不就是(batch, seq_len, hidden_size)吗？而且参数默认为False，也就是它鼓励你第一维不是batch，更奇怪了。

取pytorch官方的一个tutorial（chatbot tutorial）中的一个图：

左边是我们的常规输入（先不考虑hidden dim，每个数字代表序列中的一个词），右边是转置后，第一维成了max_length。我们知道在操作时第一维一般可视为“循环”维度，因此左边一个循环项是一个序列，无法同时经LSTM处理，而右边跨batch的循环项相当于当前time step下所有序列的当前词，可以并行过LSTM。（当然不管你是否batch_first它都是这么处理的，这个参数应该只是提醒一下这个trick）

pack&pad

（感觉说起来没那么简单，所以加了个小标题。）

前面说过nn.LSTM的输入也可以是“packed”形式，那么这是个什么形式？

先不问为什么，看一下pack和pad的操作是怎样的。

---

torch.nn.utils.rnn.pack_sequence(sequences, enforce_sorted=True)

官方API：https://pytorch.org/docs/stable/nn.html?highlight=lstm#torch.nn.utils.rnn.pack_sequence

这是pack操作，输入的sequences是tensor组成的list，要求按长度从大到小排序。官网的例子：

---

torch.nn.utils.rnn.pad_sequence(sequences, batch_first=False, padding_value=0)

官方API：https://pytorch.org/docs/stable/nn.html?highlight=lstm#torch.nn.utils.rnn.pad_sequence

这是pad操作，sequences也是list。这个比较好理解，就是给list里的tensor都用padding_value来pad成最长的长度，并组合成一个tensor：

看了这两个操作，隐隐约约和前面的LSTM联系起来了。我们知道一个batch里的序列不一定等长，需要pad操作用0把它们都填充成max_length长度。但前面说了LSTM的一次forward对应一个time step，接收的是across batches的输入，这就导致短序列可能在当前time step上已经结束，而你还是在给它输入东西（pad），这就会对结果产生影响（可以对照公式看看，即便输入全0还是会有影响）。我们想要的效果是，LSTM知道batch中每个序列的长度，等到某个序列输入结束后下面的time step就不带它了。

传统的pad不能用，LSTM需要一种其它的方法来处理变长输入。这时我们观察刚看到的pack操作，感觉终于明白了它的道理。官方的例子有点混淆，我写了一个更直观的：

把这个例子看成是LSTM处理一个batch的过程，注意看成转置的形式，即batch_first=False，也就是[4,1,9]是第一个序列，[5,2]是第二个序列…max_length=3，batch_size=5。从输出可以看出其实是一个很简单的过程，有点像稀疏矩阵的存储方法，先都塞到一起再记录位置（这里是长度）。

这两个函数都是基本操作，一般不会直接使用。常用的是下面这两个：

---

torch.nn.utils.rnn.pack_padded_sequence(input, lengths, batch_first=False, enforce_sorted=True)

官方API：https://pytorch.org/docs/stable/nn.html?highlight=lstm#torch.nn.utils.rnn.pack_padded_sequence

顾名思义，pack一个经过pad的sequence，因为我们一般在处理数据时就已经将序列pad成等长了。lengths即为序列的长度。

---

torch.nn.utils.rnn.pad_packed_sequence(sequence, batch_first=False, padding_value=0.0, total_length=None)

官方API：https://pytorch.org/docs/stable/nn.html?highlight=lstm#torch.nn.utils.rnn.pad_packed_sequence

这是上面函数的逆操作，再pad回去供后续使用。这里的total_length是个很实用的参数，在下面的例子中可以看到。

---

一个完整的例子：

import torch as t
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

a = t.tensor([[1,2,3],[6,0,0],[4,5,0]]) #(batch_size, max_length)
lengths = t.tensor([3,1,2])

# 排序
a_lengths, idx = lengths.sort(0, descending=True)
_, un_idx = t.sort(idx, dim=0)
a = a[idx]

# 定义层 
emb = t.nn.Embedding(20,2,padding_idx=0) 
lstm = t.nn.LSTM(input_size=2, hidden_size=4, batch_first=True) 

a_input = emb(a)
a_packed_input = t.nn.utils.rnn.pack_padded_sequence(input=a_input, lengths=a_lengths, batch_first=True)
packed_out, _ = lstm(a_packed_input)
out, _ = pad_packed_sequence(packed_out)
# 根据un_idx将输出转回原输入顺序
out = t.index_select(out, 0, un_idx)

上面便是常用的使用方法（个人认为完全可以封装到LSTM函数里，不知道为什么要这么设计）。但此时假设另一个batch，b：

# b是另一个batch
b = t.tensor([[7,8,0],[9,0,0],[10,0,0]])

batch中的最大长度为2，而对于整个数据流来说max_length=3，这就导致b经LSTM后pad的结果与整体的长度不匹配，此时设置pack_padded_sequence的total_length=3即可。

---

有不正确的地方还请指正！