Upload model.py

model.py

@@ -0,0 +1,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)