| import torch |
| import torch.nn as nn |
| import torchvision.models as models |
|
|
| class img_fe_class(nn.Module): |
| def __init__(self): |
| super(img_fe_class, self).__init__() |
| weights = models.ResNet50_Weights.DEFAULT |
| self.resnet = models.resnet50(weights=weights) |
| self.resnet = nn.Sequential(*list(self.resnet.children())[:-1]) |
| for param in self.resnet.parameters(): |
| param.requires_grad = False |
|
|
| def forward(self, imgs): |
| features = self.resnet(imgs) |
| return features.reshape(features.size(0), -1) |
|
|
| class text_fe_class(nn.Module): |
| def __init__(self, vocab_size, embedding_dim, hidden_size, num_layers): |
| super(text_fe_class, self).__init__() |
| self.embedding = nn.Embedding(vocab_size, embedding_dim) |
| self.rnn = nn.LSTM(embedding_dim, hidden_size, num_layers, batch_first=True) |
|
|
| def forward(self, captions, hidden): |
| embeddings = self.embedding(captions) |
| outputs, hidden = self.rnn(embeddings, hidden) |
| return outputs, hidden |
|
|
| class image_captioning_model(nn.Module): |
| def __init__(self, vocab_size, embedding_dim, hidden_size, num_layers): |
| super(image_captioning_model, self).__init__() |
| self.img_fe = img_fe_class() |
| self.text_fe = text_fe_class(vocab_size, embedding_dim, hidden_size, num_layers) |
| self.fc_out = nn.Linear(hidden_size, vocab_size) |
| self.hidden_size = hidden_size |
| self.num_layers = num_layers |
| self.img_to_hidden = nn.Linear(2048, hidden_size) |
|
|
| def forward(self, imgs, captions): |
| img_features = self.img_fe(imgs) |
| hidden = self.img_to_hidden(img_features).unsqueeze(0).repeat(self.num_layers, 1, 1) |
| cell = torch.zeros_like(hidden) |
| outputs, _ = self.text_fe(captions, (hidden, cell)) |
| outputs = self.fc_out(outputs) |
| return outputs |
|
|
| def generate_caption(self, img, tokenizer, max_len=30, device="cuda"): |
| self.eval() |
| with torch.no_grad(): |
| img = img.unsqueeze(0).to(device) |
| img_features = self.img_fe(img) |
| hidden = self.img_to_hidden(img_features).unsqueeze(0).repeat(self.num_layers, 1, 1) |
| cell = torch.zeros_like(hidden) |
| caption = [tokenizer.tok_to_ind["<BOS>"]] |
| for _ in range(max_len): |
| inputs = torch.LongTensor(caption[-1]).unsqueeze(0).unsqueeze(0).to(device) |
| outputs, (hidden, cell) = self.text_fe(inputs, (hidden, cell)) |
| outputs = self.fc_out(outputs[:, -1, :]) |
| next_token = outputs.argmax(dim=-1).item() |
| caption.append(next_token) |
| if next_token == tokenizer.tok_to_ind["<EOS>"]: |
| break |
| return tokenizer.decode(caption) |
