File size: 2,867 Bytes
ad8fd94 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | 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) # 2048 - размер выхода ResNet50
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) # Для LSTM нужно также cell state
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) |