Spaces:
Build error
Build error
| 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__() | |
| # Загружаем полную модель ResNet50 с предобученными весами | |
| self.resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) | |
| # Замораживаем все параметры | |
| for param in self.resnet.parameters(): | |
| param.requires_grad = False | |
| def forward(self, imgs): | |
| # Используем все слои до avgpool | |
| x = self.resnet.conv1(imgs) | |
| x = self.resnet.bn1(x) | |
| x = self.resnet.relu(x) | |
| x = self.resnet.maxpool(x) | |
| x = self.resnet.layer1(x) | |
| x = self.resnet.layer2(x) | |
| x = self.resnet.layer3(x) | |
| x = self.resnet.layer4(x) | |
| x = self.resnet.avgpool(x) | |
| return torch.flatten(x, 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) | |
| self.num_layers = num_layers | |
| 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.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.text_fe.num_layers, 1, 1) | |
| cell = torch.zeros_like(hidden) | |
| outputs, _ = self.text_fe(captions, (hidden, cell)) | |
| outputs = self.fc_out(outputs) | |
| return outputs |