虽然看了一些很好的blog了解了LSTM的内部机制,但对框架中的lstm输入输出和各个参数还是没有一个清晰的认识,今天打算彻底把理论和实现联系起来,再分析一下pytorch中的LSTM实现。
先说理论部分。一个非常有名的blog把原理讲得很清楚,推荐参考。总之就是这些公式:
简单来说就是,LSTM一共有三个门,输入门,遗忘门,输出门,$i,f,o$分别为三个门的程度参数,$g$是对输入的常规RNN操作。公式里可以看到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单元,只不过第一层输入是,中间层输入是。
剩下的参数就比较好理解了,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是个很实用的参数,在下面的例子中可以看到。
一个完整的例子:
1 | import torch as t |
上面便是常用的使用方法(个人认为完全可以封装到LSTM函数里,不知道为什么要这么设计)。但此时假设另一个batch,b:
1 | # b是另一个batch |
batch中的最大长度为2,而对于整个数据流来说max_length=3,这就导致b经LSTM后pad的结果与整体的长度不匹配,此时设置pack_padded_sequence的total_length=3即可。
有不正确的地方还请指正!