pytorch中LSTM的细节分析理解（转载）

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

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

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

那么实际应用时，我们并不关心细胞本身的状态，而是要拿到它呈现出的状态 $h'$ 作为最终输出。

以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$ 实际上是每个layer最后一个状态（纵向）输出的拼接。

output of shape (seq_len, batch, num_directions * hidden_size): tensor containing the output features (h_t) from the last layer of the LSTM, for each t.

h_n of shape (num_layers * num_directions, batch, hidden_size): tensor containing the hidden state for t = seq_len.

c_n of shape (num_layers * num_directions, batch, hidden_size): tensor containing the cell state for t = seq_len.

也就是说，对于单向LSTM来说， $h_n[-1,:,:]$ 就是 $output[-1,:,:]$ ，相当于序列最后一个时间步的输出。如果使用LSTM的目的是得到整个序列的embedding，与序列长度无关，由于LSTM具有序列信息传递性，因此一般可以取 $h_n[-1,:,:]$ 当作序列embedding。但双向LSTM推广后，每个时间步的隐层输出都可以作为当前词的一个融合了上下文的embedding，因此BiLSTM可以视为一种词级别的encoder方法，得到的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的处理。

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_first这个参数呢？常规的输入不就是(batch, seq_len, input_size)吗？而且参数默认为False，也就是它鼓励你第一维不是batch，更奇怪了。

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

图中batch size = 6， max_seq_length = 4。左：[Batch_size, Seq_len, 1] 右：[Seq_len, Batch_size, 1]

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

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 batch samples的输入，这就导致短序列可能在当前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个样本。从pack_sequence的输出可以看出其实是一个很简单的过程，有点像稀疏矩阵的存储方法，先都塞到一起再记录位置（这里是长度）。先把所有样本序列塞到一个张量里，然后记录每个样本序列的长度。

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

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是个很实用的参数，在下面的例子中可以看到。

a = torch.tensor([[1., 2., 3., 4, 1], [1, 2, 3, 0, 0]])

b = torch.nn.utils.rnn.pack_padded_sequence(a, torch.tensor([2, 2, 2, 1, 1]))

b

>>> PackedSequence(data=tensor([1., 2., 3., 4., 1., 1., 2., 3.]), batch_sizes=tensor([5, 3]), sorted_indices=None, unsorted_indices=None)

c = torch.nn.utils.rnn.pad_packed_sequence(b)

c

>>> (tensor([[1., 2., 3., 4., 1.], [1., 2., 3., 0., 0.]]), tensor([2, 2, 2, 1, 1]))

一个完整的例子：

 1 import torch as t
 2 from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
 3 
 4 a = t.tensor([[1,2,3],[6,0,0],[4,5,0]]) #(batch_size, max_length)
 5 lengths = t.tensor([3,1,2])
 6 
 7 # 排序
 8 a_lengths, idx = lengths.sort(0, descending=True)
 9 _, un_idx = t.sort(idx, dim=0)
10 a = a[idx]
11 
12 # 定义层 
13 emb = t.nn.Embedding(20,2,padding_idx=0) 
14 lstm = t.nn.LSTM(input_size=2, hidden_size=4, batch_first=True) 
15 
16 a_input = emb(a)
17 a_packed_input = t.nn.utils.rnn.pack_padded_sequence(input=a_input, lengths=a_lengths, batch_first=True)
18 packed_out, _ = lstm(a_packed_input)
19 out, _ = pad_packed_sequence(packed_out, batch_first=True)
20 # 根据un_idx将输出转回原输入顺序
21 out = t.index_select(out, 0, un_idx)

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

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

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

转自知乎：

https://zhuanlan.zhihu.com/p/79064602