【PyTorch学习笔记】23：nn.LSTM和nn.LSTMCell的使用

1 nn.LSTM

该模块一次构造完若干层的LSTM。

1.1 构造方法

使用nn.LSTM可以直接构建若干层的LSTM，构造时传入的三个参数和nn.RNN一样，依次是：
$[feature\_len,hidden\_len,num\_layers]$[feature_len,hidden_len,num_layers]

其中hidden_len既用于隐藏状态$h_t$ht​的尺寸，也用于记忆单元$C_t$Ct​的尺寸。实际上，两者的shape是一样的，都是$[num\_layers,batch,hidden\_len]$[num_layers,batch,hidden_len]。

1.2 forward方法

回顾一下nn.RNN的forward方法，它是：
$out,h_t = rnn(x,h_0)$out,ht​=rnn(x,h0​)

而对于nn.LSTM也是类似的，无非是要多提供初始的$C_0$C0​，并在最终多得到一个$C_t$Ct​：
$out, (h_t, C_t) = lstm(x,(h_0,C_0))$out,(ht​,Ct​)=lstm(x,(h0​,C0​))

其中，$x$x仍然是一次性将当前batch所有seq_len个时刻的样本都输入，其shape是$[seq\_len,batch,feature\_len]$[seq_len,batch,feature_len]。

在任意$t$t时刻的$h_t$ht​和$C_t$Ct​两类Memory其shape都是$[num\_layers,batch,hidden\_len]$[num_layers,batch,hidden_len]。

最终的输出$out$out同样是在所有时刻最后一层上的输出（特别注意，LSTM中的输出是$h$h不是$C$C），因此是将feature_len规约成了hidden_len那么长，所以其shape是$[seq\_len,batch,hidden\_len]$[seq_len,batch,hidden_len]。

1.3 例子

import torch
from torch import nn

# 如可以解释成:4层的LSTM,输入的每个词用100维向量表示,隐藏单元和记忆单元的尺寸是20
lstm = nn.LSTM(input_size=100, hidden_size=20, num_layers=4)

# 输入的x:其中batch是3可表示有三句话,seq_len=10表示每句话10个单词,feature_len=100表示每个单词表示为长100的向量
x = torch.randn(10, 3, 100)

# 前向计算过程,这里不传入h_0和C_0则会默认初始化
out, (h, c) = lstm(x)
print(out.shape)  # torch.Size([10, 3, 20]) 最后一层10个时刻的输出
print(h.shape)  # torch.Size([4, 3, 20]) 隐藏单元
print(c.shape)  # torch.Size([4, 3, 20]) 记忆单元

2 nn.LSTMCell

该模块构建LSTM中的一个Cell，同一层会共享这一个Cell，但要手动处理每个时刻的迭代计算过程。如果要建立多层的LSTM，就要建立多个nn.LSTMCell。

2.1 构造方法

构造方法和nn.LSTM类似，依次传入feature_len和hidden_len，因为这只是一个计算单元，所以不涉及层数。

2.2 forward方法

回顾一下nn.RNNCell的forward方法，它是：
$h_t = rnncell(x,h_{t-1})$ht​=rnncell(x,ht−1​)

即上一时刻的输出$h_{t-1}$ht−1​经nn.RNNCell前向计算得到这一时刻的$h_t$ht​。

对于nn.LSTMCell也是类似，但因为LSTM计算单元中还涉及对上一时刻的记忆$C_{t-1}$Ct−1​的使用，所以是：
$h_t,C_t = lstmcell(x_t,(h_{t-1},C_{t-1}))$ht​,Ct​=lstmcell(xt​,(ht−1​,Ct−1​))

因为输入$x_t$xt​只是t时刻的输入，不涉及seq_len，所以其shape是$[batch,feature\_len]$[batch,feature_len]。

而$h_t$ht​和$C_t$Ct​在这里只是t时刻本层的隐藏单元和记忆单元，不涉及num_layers，所以其shape是$[batch,hidden\_len]$[batch,hidden_len]。

2.3 一层的例子

每个时刻传入新的输入$x_t$xt​和上一时刻的隐藏单元$h_{t-1}$ht−1​和记忆单元$C_{t-1}$Ct−1​，并把这两个单元更新。

import torch
from torch import nn

# 一层的LSTM计算单元,输入的feature_len=100,隐藏单元和记忆单元hidden_len=20
cell = nn.LSTMCell(input_size=100, hidden_size=20)

# 初始化隐藏单元h和记忆单元C,取batch=3
h = torch.zeros(3, 20)
C = torch.zeros(3, 20)

# 这里是seq_len=10个时刻的输入,每个时刻shape都是[batch,feature_len]
xs = [torch.randn(3, 100) for _ in range(10)]

# 对每个时刻,传入输入x_t和上个时刻的h_{t-1}和C_{t-1}
for xt in xs:
    h, C = cell(xt, (h, C))

print(h.shape)  # torch.Size([3, 20])
print(C.shape)  # torch.Size([3, 20])

2.4 两层的例子

在最底下一层l0层和上面的例子一样，上层还需要接受下层的输出$h_t$ht​作为当前输入，然后同样是依赖本层上一时刻的$h$h和$C$C更新本层的$h$h和$C$C。

特别要注意DL学习笔记4中学习的LSTM单元的输入输出结构，向上层传递的是$h$h而不是$C$C。

import torch
from torch import nn

# 输入的feature_len=100,变到该层隐藏单元和记忆单元hidden_len=30
cell_l0 = nn.LSTMCell(input_size=100, hidden_size=30)
# hidden_len从l0层的30变到这一层的20
cell_l1 = nn.LSTMCell(input_size=30, hidden_size=20)

# 分别初始化l0层和l1层的隐藏单元h和记忆单元C,取batch=3
# 注意l0层的hidden_len=30
h_l0 = torch.zeros(3, 30)
C_l0 = torch.zeros(3, 30)
# 而到l1层之后hidden_len=20
h_l1 = torch.zeros(3, 20)
C_l1 = torch.zeros(3, 20)

# 这里是seq_len=10个时刻的输入,每个时刻shape都是[batch,feature_len]
xs = [torch.randn(3, 100) for _ in range(10)]

# 对每个时刻,从下到上计算本时刻的所有层
for xt in xs:
    h_l0, C_l0 = cell_l0(xt, (h_l0, C_l0))  # l0层直接接受xt输入
    h_l1, C_l1 = cell_l1(h_l0, (h_l1, C_l1))  # l1层接受l0层的输出h为输入

# 最后shape是不变的
print(h_l0.shape)  # torch.Size([3, 30])
print(C_l0.shape)  # torch.Size([3, 30])
print(h_l1.shape)  # torch.Size([3, 20])
print(C_l1.shape)  # torch.Size([3, 20])