Views
No views yet
1from torch import nn, optim
2import torch
3import os
4import pandas as pd
5import torch.nn.functional as F
6
7class CFG:
8 debug = False
9
10 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11
12 model_name = 'facebook/dino-vitb16'
13 image_embedding = 768
14 text_encoder_model = "microsoft/deberta-v3-base"
15 text_embedding = 768
16 text_tokenizer = "microsoft/deberta-v3-base"
17 max_length = 200
18
19 pretrained = True # for both image encoder and text encoder
20 trainable = True # for both image encoder and text encoder
21 temperature = 1.0
22
23 # image size
24 size = 224
25
26 # for projection head; used for both image and text encoders
27 projection_dim = 768
28 projection_width = 768
29 projection_heads = 16
30
31class ImageEncoder(nn.Module):
32 """
33 Encode images to a fixed size vector
34 """
35 def __init__(
36 self, model_name=CFG.model_name, pretrained=CFG.pretrained, trainable=CFG.trainable
37 ):
38 super().__init__()
39 #self.model = timm.create_model(
40 # model_name, pretrained, num_classes=0, global_pool="avg"
41 #)
42 self.model = ViTModel.from_pretrained(model_name)
43 self.feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
44 for p in self.model.parameters():
45 p.requires_grad = trainable
46
47 def forward(self, x):
48 return self.model(x).last_hidden_state[:,0,:]
49
50class TextEncoder(nn.Module):
51 def __init__(self, model_name=CFG.text_encoder_model, pretrained=CFG.pretrained, trainable=CFG.trainable):
52 super().__init__()
53 if pretrained:
54 self.model = DebertaV2Model.from_pretrained(model_name)
55 else:
56 self.model = DebertaV2Model(config=DistilBertConfig())
57
58 for p in self.model.parameters():
59 p.requires_grad = trainable
60
61 # we are using the CLS token hidden representation as the sentence's embedding
62 self.target_token_idx = 0
63
64 def forward(self, input_ids, attention_mask):
65 output = self.model(input_ids=input_ids, attention_mask=attention_mask)
66 last_hidden_state = output.last_hidden_state
67 return last_hidden_state[:, self.target_token_idx, :]
68
69from collections import OrderedDict
70from torch import nn
71import torch
72
73class QuickGELU(nn.Module):
74 def forward(self, x: torch.Tensor):
75 return x * torch.sigmoid(1.702 * x)
76
77class ProjectionHead(nn.Module):
78 def __init__(
79 self,
80 embedding_dim,
81 projection_dim=384,
82 dropout=0.1
83 ):
84 super().__init__()
85 self.projection = nn.Linear(embedding_dim, projection_dim)
86 self.gelu = QuickGELU()
87 self.fc = nn.Linear(projection_dim, projection_dim)
88 self.dropout = nn.Dropout(dropout)
89 self.layer_norm = nn.LayerNorm(projection_dim)
90
91 def forward(self, x):
92 projected = self.projection(x)
93 x = self.gelu(projected)
94 x = self.fc(x)
95 x = self.dropout(x)
96 x = x + projected
97 x = self.layer_norm(x)
98 return x
99
100
101class CLIPModel(nn.Module):
102 def __init__(self):
103 super().__init__()
104 self.feature_extractor = ViTFeatureExtractor.from_pretrained('facebook/dino-vitb16')
105 self.image_encoder = ViTModel.from_pretrained('facebook/dino-vitb16').to('cuda')
106 self.text_encoder = DebertaV2Model.from_pretrained("microsoft/deberta-v3-base").to('cuda')
107 self.image_projection = ProjectionHead(768).to('cuda')
108 self.text_projection = ProjectionHead(768).to('cuda')
109 self.temperature = 1.0
110 def encode_image(self, x):
111 x = self.image_encoder(x).last_hidden_state[:,0,:]
112 x = self.image_projection(x)
113 return x
114 def encode_text(self, x):
115 x = self.text_encoder(x['input_ids'].cuda()).last_hidden_state[:,0,:]
116 x = self.text_projection(x)
117 return x
118 def tokenize(self, x):
119 return tokenizer(x, return_tensors='pt', padding=True, truncation=True)
120 def preprocess(self, x):
121 x = self.feature_extractor(x, return_tensors='pt')
122 return x.convert_to_tensors()['pixel_values'].to('cuda')
123
124def cross_entropy(preds, targets, reduction='none'):
125 log_softmax = nn.LogSoftmax(dim=-1)
126 loss = (-targets * log_softmax(preds)).sum(1)
127 if reduction == "none":
128 return loss
129 elif reduction == "mean":
130 return loss.mean()1from huggingface_hub import hf_hub_url, cached_download
2model_name = "model.pt"
3model_url = hf_hub_url("crumb/pmclip-test-run-checkpoints-10", filename=model_name)
4file_path = cached_download(model_url)
5tokenizer = DebertaV2Tokenizer.from_pretrained(CFG.text_tokenizer)
6model = CLIPModel().to(CFG.device)
7model.load_state_dict(torch.load(file_path))
8model = model.eval().cuda()
9print("num params", sum(p.numel() for p in model.parameters()) // 10000 / 100, 'M')
10# num params 271.1 M