Max Pooling over 2D input?

Hey guys!

I’m currently working on an LSTM for text classification. I’d like to build a max pooling layer in the model, but I’m only working with 2D inputs. Is there a way to do a max pooling after getting lstm_out? I want all the listings of the first dimension to be pooled, so I get the most important feature out of it. So self.hidden2label should only get one feature value for choosing the right label. Here is how my model looks like:

class LSTM(nn.Module):

    def __init__(self, embedding_dim, hidden_dim, alphabet_size, label_size, batch_size):
        super(LSTM, self).__init__()
        self.hidden_dim = hidden_dim
        self.batch_size = batch_size

        self.word_embeddings = nn.Embedding(alphabet_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, bidirectional=True)
        self.hidden2label = nn.Linear(hidden_dim * 2, label_size) # Ä: hidden_dim * 2 for bidirectional
        print(self.hidden2label)
        self.hidden = self.init_hidden()

    def init_hidden(self):

        h0 = Variable(torch.zeros(1 * 2, self.batch_size, self.hidden_dim))
        c0 = Variable(torch.zeros(1 * 2, self.batch_size, self.hidden_dim))
        return h0, c0

    def forward(self, sentence):
        embeds = self.word_embeddings(sentence)
        x = embeds.view(len(sentence), self.batch_size, -1)
        lstm_out, self.hidden = self.lstm(x, self.hidden)

        sent_rep = torch.cat((lstm_out[-1, :, :self.hidden_dim], lstm_out[0, :, self.hidden_dim:]), 1)

        y = self.hidden2label(sent_rep)

        return y

I hope my goal is comprehensible. Thank you for helping! :slight_smile: