带有顺序模块的PyTorch中的简单LSTM

Question

在PyTorch中,我们可以通过多种方式定义体系结构.在这里,我想使用该Sequential模块创建一个简单的LSTM网络.

在Lua的火炬中,我通常会选择:

model = nn.Sequential()
model:add(nn.SplitTable(1,2))
model:add(nn.Sequencer(nn.LSTM(inputSize, hiddenSize)))
model:add(nn.SelectTable(-1)) -- last step of output sequence
model:add(nn.Linear(hiddenSize, classes_n))

但是,在PyTorch中,我找不到相当于SelectTable获得最后一个输出.

nn.Sequential(
  nn.LSTM(inputSize, hiddenSize, 1, batch_first=True),
  # what to put here to retrieve last output of LSTM ?,
  nn.Linear(hiddenSize, classe_n))

Answer 1

定义一个类来提取最后一个单元格的输出：

# LSTM() returns tuple of (tensor, (recurrent state))
class extract_tensor(nn.Module):
    def forward(self,x):
        # Output shape (batch, features, hidden)
        tensor, _ = x
        # Reshape shape (batch, hidden)
        return tensor[:, -1, :]

nn.Sequential(
    nn.LSTM(inputSize, hiddenSize, 1, batch_first=True),
    extract_tensor(),
    nn.Linear(hiddenSize, classe_n)
)

Answer 2

根据LSTM 单元文档，输出参数的形状为 (seq_len、batch、hidden_​​size * num_directions)，因此您可以通过以下方式轻松获取序列的最后一个元素：

rnn = nn.LSTM(10, 20, 2)
input = Variable(torch.randn(5, 3, 10)) 
h0 = Variable(torch.randn(2, 3, 20))
c0 = Variable(torch.randn(2, 3, 20))
output, hn = rnn(input, (h0, c0))
print(output[-1]) # last element

PyTorch 中的张量操作和神经网络设计比 Torch 中的要容易得多，因此您很少需要使用容器。事实上，正如针对前 Torch 用户的教程 PyTorch中所述，PyTorch 是围绕 Autograd 构建的，因此您不再需要担心容器。但是，如果您想使用旧的 Lua Torch 代码，您可以查看Legacy 包。