|
import torch |
|
from torch import nn |
|
import timm |
|
from transformers import DistilBertModel, DistilBertConfig |
|
import config as CFG |
|
|
|
|
|
class ImageEncoder(nn.Module): |
|
""" |
|
Encode images to a fixed size vector |
|
""" |
|
|
|
def __init__( |
|
self, model_name=CFG.model_name, pretrained=CFG.pretrained, trainable=CFG.trainable |
|
): |
|
super().__init__() |
|
self.model = timm.create_model( |
|
model_name, pretrained, num_classes=0, global_pool="avg" |
|
) |
|
for p in self.model.parameters(): |
|
p.requires_grad = trainable |
|
|
|
def forward(self, x): |
|
return self.model(x) |
|
|
|
|
|
class TextEncoder(nn.Module): |
|
def __init__(self, model_name=CFG.text_encoder_model, pretrained=CFG.pretrained, trainable=CFG.trainable): |
|
super().__init__() |
|
if pretrained: |
|
self.model = DistilBertModel.from_pretrained(model_name) |
|
else: |
|
self.model = DistilBertModel(config=DistilBertConfig()) |
|
|
|
for p in self.model.parameters(): |
|
p.requires_grad = trainable |
|
|
|
|
|
self.target_token_idx = 0 |
|
|
|
def forward(self, input_ids, attention_mask): |
|
output = self.model(input_ids=input_ids, attention_mask=attention_mask) |
|
last_hidden_state = output.last_hidden_state |
|
return last_hidden_state[:, self.target_token_idx, :] |
|
|
|
|
|
|
|
class ProjectionHead(nn.Module): |
|
def __init__( |
|
self, |
|
embedding_dim, |
|
projection_dim=CFG.projection_dim, |
|
dropout=CFG.dropout |
|
): |
|
super().__init__() |
|
self.projection = nn.Linear(embedding_dim, projection_dim) |
|
self.gelu = nn.GELU() |
|
self.fc = nn.Linear(projection_dim, projection_dim) |
|
self.dropout = nn.Dropout(dropout) |
|
self.layer_norm = nn.LayerNorm(projection_dim) |
|
|
|
def forward(self, x): |
|
projected = self.projection(x) |
|
x = self.gelu(projected) |
|
x = self.fc(x) |
|
x = self.dropout(x) |
|
x = x + projected |
|
x = self.layer_norm(x) |
|
return x |
|
|
|
|
