Simple LSTM in PyTorch with Sequential module

user_6246880 0
Viewed 8550

In PyTorch, we can define architectures in multiple ways. Here, I'd like to create a simple LSTM network using the Sequential module.

In Lua's torch I would usually go with:

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))

However, in PyTorch, I don't find the equivalent of SelectTable to get the last output.

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))
torch recurrent-neural-network lstm pytorch deep-learning
3 Answers

Define a class to extract the last cell output:

# 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)
)
user_25808310

As far as I'm concerned there's nothing like a SplitTable or a SelectTable in PyTorch. That said, you are allowed to concatenate an arbitrary number of modules or blocks within a single architecture, and you can use this property to retrieve the output of a certain layer. Let's make it more clear with a simple example.

Suppose I want to build a simple two-layer MLP and retrieve the output of each layer. I can build a custom class inheriting from nn.Module:

class MyMLP(nn.Module):

    def __init__(self, in_channels, out_channels_1, out_channels_2):
        # first of all, calling base class constructor
        super().__init__()
        # now I can build my modular network
        self.block1 = nn.Linear(in_channels, out_channels_1)
        self.block2 = nn.Linear(out_channels_1, out_channels_2)

    # you MUST implement a forward(input) method whenever inheriting from nn.Module
    def forward(x):
        # first_out will now be your output of the first block
        first_out = self.block1(x)
        x = self.block2(first_out)
        # by returning both x and first_out, you can now access the first layer's output
        return x, first_out

In your main file you can now declare the custom architecture and use it:

from myFile import MyMLP
import numpy as np

in_ch = out_ch_1 = out_ch_2 = 64
# some fake input instance
x = np.random.rand(in_ch)

my_mlp = MyMLP(in_ch, out_ch_1, out_ch_2)
# get your outputs
final_out, first_layer_out = my_mlp(x)

Moreover, you could concatenate two MyMLP in a more complex model definition and retrieve the output of each one in a similar way. I hope this is enough to clarify, but in case you have more questions, please feel free to ask, since I may have omitted something.

user_98850030
Related