Spaces:
Build error
Build error
File size: 2,286 Bytes
3462329 15012bd 3462329 15012bd 3462329 15012bd 3462329 15012bd 3462329 15012bd 3462329 15012bd | 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 | 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 |