| import torch.nn as nn
|
| import torch
|
| import numpy as np
|
| from model.encoder import ImageEncoder, RobertaEncoder
|
| import torch.nn.functional as F
|
| class LVL(nn.Module):
|
| def __init__(self):
|
| super(LVL, self).__init__()
|
| self.image_encoder = ImageEncoder()
|
| self.text_encoder = RobertaEncoder()
|
| self.t_prime = nn.Parameter(torch.ones([]) * np.log(0.07))
|
| self.b = nn.Parameter(torch.ones([]) * 0)
|
|
|
| def get_images_features(self,images):
|
| image_embeddings = self.image_encoder(images)
|
| image_embeddings = F.normalize(image_embeddings, p=2, dim=-1)
|
| return image_embeddings
|
|
|
| def get_texts_feature(self,input_ids,attention_mask):
|
| text_embeddings = self.text_encoder(input_ids, attention_mask)
|
| text_embeddings = F.normalize(text_embeddings, p=2, dim=-1)
|
| return text_embeddings
|
|
|
| def forward(self, images, input_ids, attention_mask):
|
| """
|
| Args:
|
| images: Tensor of shape (batch_size, 3, 224, 224)
|
| input_ids: Tensor of shape (batch_size, seq_length)
|
| attention_mask: Tensor of shape (batch_size, seq_length)
|
|
|
| Returns:
|
| Image and text embeddings normalized for similarity calculation
|
| """
|
|
|
| image_embeddings = self.get_images_features(images)
|
| text_embeddings = self.get_texts_feature(input_ids, attention_mask)
|
| return image_embeddings, text_embeddings
|
|
|
