I’m trying to add a Attention Mechanism to a Decoder for the purpose of improving a Image Captioning Model.
I’m using this tutorial: Github and I’m trying to add this attention mechanism: Github.
The problem is that it seems like the shapes and sizes of the tensors don’t match: Shape of features: torch.Size([128, 256]) Shape of hiddens: torch.Size([128, 1, 21, 512])

I’m trying to reshape or resize the tensors so they match via PyTorch .resize and .reshape, i’ve tried .unsqueeze & squeeze too but they don’t change the shapes, when I do resize or reshape it appears this error:

# when I do:
        new_hiddens = hiddens.reshape(128, 1, 23, 256)
# it says:
RuntimeError: shape '[128, 1, 23, 256]' is invalid for input of size 1441792
#and when I do:
        new_hiddens = hiddens.resize(128, 256)
#it says:
requested resize to 128x256 (32768 elements in total), but the given tensor has a size of 128x1x26x512 (1703936 elements). autograd's resize can only change the shape of a given tensor, while preserving the number of elements. 

Then I ask GPT and it says that maybe it’s not because of the tensors shapes but how are they used, and it makes sense because they are from 2 different examples.
So I hope somebody more experienced than me can help me identify where in my Attention mechanism it’s expecting a different type of tensors.

class Attention(nn.Module):
    def __init__(self, encoder_dim, decoder_dim, attention_dim):
        super(Attention, self).__init__()
        self.encoder_att = nn.Linear(encoder_dim, attention_dim)
        self.decoder_att = nn.Linear(decoder_dim, attention_dim)
        self.full_att = nn.Linear(attention_dim, 1)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=2)

    def forward(self, encoder_out, decoder_hidden):
        att1 = self.encoder_att(encoder_out)  # (batch_size, 1, attention_dim)
        att2 = self.decoder_att(decoder_hidden)  # (batch_size, seq_len, attention_dim)
        att = self.full_att(self.relu(att1 + att2)).squeeze(2)  # (batch_size, seq_len)
        alpha = self.softmax(att)  # (batch_size, seq_len)
        attention_weighted_encoding = (encoder_out.unsqueeze(1) * alpha.unsqueeze(2)).sum(dim=1)  # (batch_size, encoder_dim)

        return attention_weighted_encoding, alpha


class DecoderRNN(nn.Module):
    def __init__(self, embed_size, hidden_size, vocab_size, num_layers, max_seq_length=20):
        super(DecoderRNN, self).__init__()
        self.embed = nn.Embedding(vocab_size, embed_size)
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)  # change here
        self.linear = nn.Linear(hidden_size, vocab_size)
        self.max_seg_length = max_seq_length
        self.attention = Attention(hidden_size, hidden_size, hidden_size)  # add attention here

    def forward(self, features, captions, lengths):
        embeddings = self.embed(captions)
        hiddens, _ = self.lstm(embeddings)
        hiddens = hiddens.unsqueeze(1)
        #new_hiddens = hiddens.resize(128, 256)
        #print("Shape of new hiddens: ", new_hiddens.shape)
        print("Shape of features: ", features.shape)
        print("Shape of hiddens: ", hiddens.shape)
        attn_weights = self.attention(features, hiddens)
        context = attn_weights.bmm(features.unsqueeze(1))  # (b, 1, n)
        hiddens = hiddens + context
        outputs = self.linear(hiddens.squeeze(1))
        return outputs