I'm a beginner in Machine Learning and I'm working with the Flickr8k dataset (it contains ~8000 images, every image has 5 captions: ~40000 pairs). I splitted the dataset in training (70%) and validation/test (15%). I trained the model (for image captioning) on my GPU for 40 epochs (I need ~400s for each epoch).
During the model training the accuracy (n° of correct words predicted) increases but if I try to predict a caption given an image, the caption generated is always like "<start> <start> <start> <start> ..." (so my LSTM predicts always the same word/token). I think maybe the issue is the implementation of my "predict" function (since that the forward pass of the model returns a better caption each training epoch). Here's some code:
class CNNEncoder(nn.Module): def __init__(self, embed_size): # embed_size is the output size of the CNN encoder super(CNNEncoder, self).__init__() self.inception = models.inception_v3(weights=Inception_V3_Weights.DEFAULT) self.inception.fc = nn.Linear(self.inception.fc.in_features, embed_size) for param in self.inception.parameters(): if param is not self.inception.fc.weight and param is not self.inception.fc.bias: param.requires_grad_(False) self.relu = nn.ReLU() self.dropout = nn.Dropout(0.5) def forward(self, input): if self.training: features, _ = self.inception(input) else: features = self.inception(input) return self.dropout(self.relu(features)) class DecoderRNN(nn.Module): def __init__(self, embed_size, hidden_size, vocab_size, n_lstm_layers=2): super(DecoderRNN, self).__init__() self.embed = nn.Embedding(vocab_size, embed_size) self.lstm = nn.LSTM(embed_size, hidden_size, num_layers=n_lstm_layers, batch_first=True) self.linear = nn.Linear(hidden_size, vocab_size) self.dropout = nn.Dropout(0.5) def forward(self, features, captions): embeddings = self.dropout(self.embed(captions)) embeddings = torch.cat((features.unsqueeze(1), embeddings), dim=1) # shape: (batch_size, caption_length + 1, embed_size) outputs, _ = self.lstm(embeddings) return self.linear(outputs) class ImageCaptioner(nn.Module): def __init__(self, embed_size, hidden_size, vocab_size, n_lstm_layers=2): super(ImageCaptioner, self).__init__() self.cnn = CNNEncoder(embed_size) self.decoderRNN = DecoderRNN(embed_size, hidden_size, vocab_size, n_lstm_layers) def forward(self, images, captions): features = self.cnn(images) # shape: (batch_size, embed_size) outputs = self.decoderRNN(features, captions) # shape: (batch_size, caption_length, vocab_size) return outputs def predict(self, images, device, vocab, max_length=30): result = torch.zeros((images.shape, max_length, len(vocab)), dtype=torch.float32, device=device) feature_vector = self.cnn(images).unsqueeze(1) states = None for i in range(max_length): outputs, states = self.decoderRNN.lstm(feature_vector, states) outputs = self.decoderRNN.linear(outputs.squeeze(1)) predicted = outputs.argmax(1) # shape: (batch_size) result[:, i, predicted] = 1 feature_vector = self.decoderRNN.embed(predicted.unsqueeze(1)) # shape: (batch_size, 1, embed_size) if (predicted == vocab["<end>"]).all(): break break return result
I already tried to change embed_size(256)/hidden_size(128)/n_lstm_layers(1) params of my model. What could be going wrong? Why does my LSTM during testing return always the same output? I already tried to test the model using the training set instead of the test set (and it always predicts only <start> tokens).