Spaces:
Running
Running
Upload models/blip_pretrain.py
Browse files- models/blip_pretrain.py +339 -0
models/blip_pretrain.py
ADDED
@@ -0,0 +1,339 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
'''
|
2 |
+
* Copyright (c) 2022, salesforce.com, inc.
|
3 |
+
* All rights reserved.
|
4 |
+
* SPDX-License-Identifier: BSD-3-Clause
|
5 |
+
* For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
|
6 |
+
* By Junnan Li
|
7 |
+
'''
|
8 |
+
from models.med import BertConfig, BertModel, BertLMHeadModel
|
9 |
+
from transformers import BertTokenizer
|
10 |
+
import transformers
|
11 |
+
transformers.logging.set_verbosity_error()
|
12 |
+
|
13 |
+
import torch
|
14 |
+
from torch import nn
|
15 |
+
import torch.nn.functional as F
|
16 |
+
|
17 |
+
from models.blip import create_vit, init_tokenizer, load_checkpoint
|
18 |
+
|
19 |
+
class BLIP_Pretrain(nn.Module):
|
20 |
+
def __init__(self,
|
21 |
+
med_config = 'configs/bert_config.json',
|
22 |
+
image_size = 224,
|
23 |
+
vit = 'base',
|
24 |
+
vit_grad_ckpt = False,
|
25 |
+
vit_ckpt_layer = 0,
|
26 |
+
embed_dim = 256,
|
27 |
+
queue_size = 57600,
|
28 |
+
momentum = 0.995,
|
29 |
+
):
|
30 |
+
"""
|
31 |
+
Args:
|
32 |
+
med_config (str): path for the mixture of encoder-decoder model's configuration file
|
33 |
+
image_size (int): input image size
|
34 |
+
vit (str): model size of vision transformer
|
35 |
+
"""
|
36 |
+
super().__init__()
|
37 |
+
|
38 |
+
self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer, 0)
|
39 |
+
|
40 |
+
if vit=='base':
|
41 |
+
checkpoint = torch.hub.load_state_dict_from_url(
|
42 |
+
url="https://dl.fbaipublicfiles.com/deit/deit_base_patch16_224-b5f2ef4d.pth",
|
43 |
+
map_location="cpu", check_hash=True)
|
44 |
+
state_dict = checkpoint["model"]
|
45 |
+
msg = self.visual_encoder.load_state_dict(state_dict,strict=False)
|
46 |
+
elif vit=='large':
|
47 |
+
from timm.models.helpers import load_custom_pretrained
|
48 |
+
from timm.models.vision_transformer import default_cfgs
|
49 |
+
load_custom_pretrained(self.visual_encoder,default_cfgs['vit_large_patch16_224_in21k'])
|
50 |
+
|
51 |
+
self.tokenizer = init_tokenizer()
|
52 |
+
encoder_config = BertConfig.from_json_file(med_config)
|
53 |
+
encoder_config.encoder_width = vision_width
|
54 |
+
self.text_encoder = BertModel.from_pretrained('bert-base-uncased',config=encoder_config, add_pooling_layer=False)
|
55 |
+
self.text_encoder.resize_token_embeddings(len(self.tokenizer))
|
56 |
+
|
57 |
+
text_width = self.text_encoder.config.hidden_size
|
58 |
+
|
59 |
+
self.vision_proj = nn.Linear(vision_width, embed_dim)
|
60 |
+
self.text_proj = nn.Linear(text_width, embed_dim)
|
61 |
+
|
62 |
+
self.itm_head = nn.Linear(text_width, 2)
|
63 |
+
|
64 |
+
# create momentum encoders
|
65 |
+
self.visual_encoder_m, vision_width = create_vit(vit,image_size)
|
66 |
+
self.vision_proj_m = nn.Linear(vision_width, embed_dim)
|
67 |
+
self.text_encoder_m = BertModel(config=encoder_config, add_pooling_layer=False)
|
68 |
+
self.text_proj_m = nn.Linear(text_width, embed_dim)
|
69 |
+
|
70 |
+
self.model_pairs = [[self.visual_encoder,self.visual_encoder_m],
|
71 |
+
[self.vision_proj,self.vision_proj_m],
|
72 |
+
[self.text_encoder,self.text_encoder_m],
|
73 |
+
[self.text_proj,self.text_proj_m],
|
74 |
+
]
|
75 |
+
self.copy_params()
|
76 |
+
|
77 |
+
# create the queue
|
78 |
+
self.register_buffer("image_queue", torch.randn(embed_dim, queue_size))
|
79 |
+
self.register_buffer("text_queue", torch.randn(embed_dim, queue_size))
|
80 |
+
self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
|
81 |
+
|
82 |
+
self.image_queue = nn.functional.normalize(self.image_queue, dim=0)
|
83 |
+
self.text_queue = nn.functional.normalize(self.text_queue, dim=0)
|
84 |
+
|
85 |
+
self.queue_size = queue_size
|
86 |
+
self.momentum = momentum
|
87 |
+
self.temp = nn.Parameter(0.07*torch.ones([]))
|
88 |
+
|
89 |
+
# create the decoder
|
90 |
+
decoder_config = BertConfig.from_json_file(med_config)
|
91 |
+
decoder_config.encoder_width = vision_width
|
92 |
+
self.text_decoder = BertLMHeadModel.from_pretrained('bert-base-uncased',config=decoder_config)
|
93 |
+
self.text_decoder.resize_token_embeddings(len(self.tokenizer))
|
94 |
+
tie_encoder_decoder_weights(self.text_decoder.bert,self.text_encoder,'','/attention')
|
95 |
+
|
96 |
+
|
97 |
+
def forward(self, image, caption, alpha):
|
98 |
+
with torch.no_grad():
|
99 |
+
self.temp.clamp_(0.001,0.5)
|
100 |
+
|
101 |
+
image_embeds = self.visual_encoder(image)
|
102 |
+
image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(image.device)
|
103 |
+
image_feat = F.normalize(self.vision_proj(image_embeds[:,0,:]),dim=-1)
|
104 |
+
|
105 |
+
text = self.tokenizer(caption, padding='max_length', truncation=True, max_length=30,
|
106 |
+
return_tensors="pt").to(image.device)
|
107 |
+
text_output = self.text_encoder(text.input_ids, attention_mask = text.attention_mask,
|
108 |
+
return_dict = True, mode = 'text')
|
109 |
+
text_feat = F.normalize(self.text_proj(text_output.last_hidden_state[:,0,:]),dim=-1)
|
110 |
+
|
111 |
+
# get momentum features
|
112 |
+
with torch.no_grad():
|
113 |
+
self._momentum_update()
|
114 |
+
image_embeds_m = self.visual_encoder_m(image)
|
115 |
+
image_feat_m = F.normalize(self.vision_proj_m(image_embeds_m[:,0,:]),dim=-1)
|
116 |
+
image_feat_all = torch.cat([image_feat_m.t(),self.image_queue.clone().detach()],dim=1)
|
117 |
+
|
118 |
+
text_output_m = self.text_encoder_m(text.input_ids, attention_mask = text.attention_mask,
|
119 |
+
return_dict = True, mode = 'text')
|
120 |
+
text_feat_m = F.normalize(self.text_proj_m(text_output_m.last_hidden_state[:,0,:]),dim=-1)
|
121 |
+
text_feat_all = torch.cat([text_feat_m.t(),self.text_queue.clone().detach()],dim=1)
|
122 |
+
|
123 |
+
sim_i2t_m = image_feat_m @ text_feat_all / self.temp
|
124 |
+
sim_t2i_m = text_feat_m @ image_feat_all / self.temp
|
125 |
+
|
126 |
+
sim_targets = torch.zeros(sim_i2t_m.size()).to(image.device)
|
127 |
+
sim_targets.fill_diagonal_(1)
|
128 |
+
|
129 |
+
sim_i2t_targets = alpha * F.softmax(sim_i2t_m, dim=1) + (1 - alpha) * sim_targets
|
130 |
+
sim_t2i_targets = alpha * F.softmax(sim_t2i_m, dim=1) + (1 - alpha) * sim_targets
|
131 |
+
|
132 |
+
sim_i2t = image_feat @ text_feat_all / self.temp
|
133 |
+
sim_t2i = text_feat @ image_feat_all / self.temp
|
134 |
+
|
135 |
+
loss_i2t = -torch.sum(F.log_softmax(sim_i2t, dim=1)*sim_i2t_targets,dim=1).mean()
|
136 |
+
loss_t2i = -torch.sum(F.log_softmax(sim_t2i, dim=1)*sim_t2i_targets,dim=1).mean()
|
137 |
+
|
138 |
+
loss_ita = (loss_i2t+loss_t2i)/2
|
139 |
+
|
140 |
+
self._dequeue_and_enqueue(image_feat_m, text_feat_m)
|
141 |
+
|
142 |
+
###============== Image-text Matching ===================###
|
143 |
+
encoder_input_ids = text.input_ids.clone()
|
144 |
+
encoder_input_ids[:,0] = self.tokenizer.enc_token_id
|
145 |
+
|
146 |
+
# forward the positve image-text pair
|
147 |
+
bs = image.size(0)
|
148 |
+
output_pos = self.text_encoder(encoder_input_ids,
|
149 |
+
attention_mask = text.attention_mask,
|
150 |
+
encoder_hidden_states = image_embeds,
|
151 |
+
encoder_attention_mask = image_atts,
|
152 |
+
return_dict = True,
|
153 |
+
)
|
154 |
+
with torch.no_grad():
|
155 |
+
weights_t2i = F.softmax(sim_t2i[:,:bs],dim=1)+1e-4
|
156 |
+
weights_t2i.fill_diagonal_(0)
|
157 |
+
weights_i2t = F.softmax(sim_i2t[:,:bs],dim=1)+1e-4
|
158 |
+
weights_i2t.fill_diagonal_(0)
|
159 |
+
|
160 |
+
# select a negative image for each text
|
161 |
+
image_embeds_neg = []
|
162 |
+
for b in range(bs):
|
163 |
+
neg_idx = torch.multinomial(weights_t2i[b], 1).item()
|
164 |
+
image_embeds_neg.append(image_embeds[neg_idx])
|
165 |
+
image_embeds_neg = torch.stack(image_embeds_neg,dim=0)
|
166 |
+
|
167 |
+
# select a negative text for each image
|
168 |
+
text_ids_neg = []
|
169 |
+
text_atts_neg = []
|
170 |
+
for b in range(bs):
|
171 |
+
neg_idx = torch.multinomial(weights_i2t[b], 1).item()
|
172 |
+
text_ids_neg.append(encoder_input_ids[neg_idx])
|
173 |
+
text_atts_neg.append(text.attention_mask[neg_idx])
|
174 |
+
|
175 |
+
text_ids_neg = torch.stack(text_ids_neg,dim=0)
|
176 |
+
text_atts_neg = torch.stack(text_atts_neg,dim=0)
|
177 |
+
|
178 |
+
text_ids_all = torch.cat([encoder_input_ids, text_ids_neg],dim=0)
|
179 |
+
text_atts_all = torch.cat([text.attention_mask, text_atts_neg],dim=0)
|
180 |
+
|
181 |
+
image_embeds_all = torch.cat([image_embeds_neg,image_embeds],dim=0)
|
182 |
+
image_atts_all = torch.cat([image_atts,image_atts],dim=0)
|
183 |
+
|
184 |
+
output_neg = self.text_encoder(text_ids_all,
|
185 |
+
attention_mask = text_atts_all,
|
186 |
+
encoder_hidden_states = image_embeds_all,
|
187 |
+
encoder_attention_mask = image_atts_all,
|
188 |
+
return_dict = True,
|
189 |
+
)
|
190 |
+
|
191 |
+
vl_embeddings = torch.cat([output_pos.last_hidden_state[:,0,:], output_neg.last_hidden_state[:,0,:]],dim=0)
|
192 |
+
vl_output = self.itm_head(vl_embeddings)
|
193 |
+
|
194 |
+
itm_labels = torch.cat([torch.ones(bs,dtype=torch.long),torch.zeros(2*bs,dtype=torch.long)],
|
195 |
+
dim=0).to(image.device)
|
196 |
+
loss_itm = F.cross_entropy(vl_output, itm_labels)
|
197 |
+
|
198 |
+
##================= LM ========================##
|
199 |
+
decoder_input_ids = text.input_ids.clone()
|
200 |
+
decoder_input_ids[:,0] = self.tokenizer.bos_token_id
|
201 |
+
decoder_targets = decoder_input_ids.masked_fill(decoder_input_ids == self.tokenizer.pad_token_id, -100)
|
202 |
+
|
203 |
+
decoder_output = self.text_decoder(decoder_input_ids,
|
204 |
+
attention_mask = text.attention_mask,
|
205 |
+
encoder_hidden_states = image_embeds,
|
206 |
+
encoder_attention_mask = image_atts,
|
207 |
+
labels = decoder_targets,
|
208 |
+
return_dict = True,
|
209 |
+
)
|
210 |
+
|
211 |
+
loss_lm = decoder_output.loss
|
212 |
+
return loss_ita, loss_itm, loss_lm
|
213 |
+
|
214 |
+
|
215 |
+
|
216 |
+
@torch.no_grad()
|
217 |
+
def copy_params(self):
|
218 |
+
for model_pair in self.model_pairs:
|
219 |
+
for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
|
220 |
+
param_m.data.copy_(param.data) # initialize
|
221 |
+
param_m.requires_grad = False # not update by gradient
|
222 |
+
|
223 |
+
|
224 |
+
@torch.no_grad()
|
225 |
+
def _momentum_update(self):
|
226 |
+
for model_pair in self.model_pairs:
|
227 |
+
for param, param_m in zip(model_pair[0].parameters(), model_pair[1].parameters()):
|
228 |
+
param_m.data = param_m.data * self.momentum + param.data * (1. - self.momentum)
|
229 |
+
|
230 |
+
|
231 |
+
@torch.no_grad()
|
232 |
+
def _dequeue_and_enqueue(self, image_feat, text_feat):
|
233 |
+
# gather keys before updating queue
|
234 |
+
image_feats = concat_all_gather(image_feat)
|
235 |
+
text_feats = concat_all_gather(text_feat)
|
236 |
+
|
237 |
+
batch_size = image_feats.shape[0]
|
238 |
+
|
239 |
+
ptr = int(self.queue_ptr)
|
240 |
+
assert self.queue_size % batch_size == 0 # for simplicity
|
241 |
+
|
242 |
+
# replace the keys at ptr (dequeue and enqueue)
|
243 |
+
self.image_queue[:, ptr:ptr + batch_size] = image_feats.T
|
244 |
+
self.text_queue[:, ptr:ptr + batch_size] = text_feats.T
|
245 |
+
ptr = (ptr + batch_size) % self.queue_size # move pointer
|
246 |
+
|
247 |
+
self.queue_ptr[0] = ptr
|
248 |
+
|
249 |
+
|
250 |
+
def blip_pretrain(**kwargs):
|
251 |
+
model = BLIP_Pretrain(**kwargs)
|
252 |
+
return model
|
253 |
+
|
254 |
+
|
255 |
+
@torch.no_grad()
|
256 |
+
def concat_all_gather(tensor):
|
257 |
+
"""
|
258 |
+
Performs all_gather operation on the provided tensors.
|
259 |
+
*** Warning ***: torch.distributed.all_gather has no gradient.
|
260 |
+
"""
|
261 |
+
tensors_gather = [torch.ones_like(tensor)
|
262 |
+
for _ in range(torch.distributed.get_world_size())]
|
263 |
+
torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
|
264 |
+
|
265 |
+
output = torch.cat(tensors_gather, dim=0)
|
266 |
+
return output
|
267 |
+
|
268 |
+
|
269 |
+
from typing import List
|
270 |
+
def tie_encoder_decoder_weights(encoder: nn.Module, decoder: nn.Module, base_model_prefix: str, skip_key:str):
|
271 |
+
uninitialized_encoder_weights: List[str] = []
|
272 |
+
if decoder.__class__ != encoder.__class__:
|
273 |
+
logger.info(
|
274 |
+
f"{decoder.__class__} and {encoder.__class__} are not equal. In this case make sure that all encoder weights are correctly initialized."
|
275 |
+
)
|
276 |
+
|
277 |
+
def tie_encoder_to_decoder_recursively(
|
278 |
+
decoder_pointer: nn.Module,
|
279 |
+
encoder_pointer: nn.Module,
|
280 |
+
module_name: str,
|
281 |
+
uninitialized_encoder_weights: List[str],
|
282 |
+
skip_key: str,
|
283 |
+
depth=0,
|
284 |
+
):
|
285 |
+
assert isinstance(decoder_pointer, nn.Module) and isinstance(
|
286 |
+
encoder_pointer, nn.Module
|
287 |
+
), f"{decoder_pointer} and {encoder_pointer} have to be of type torch.nn.Module"
|
288 |
+
if hasattr(decoder_pointer, "weight") and skip_key not in module_name:
|
289 |
+
assert hasattr(encoder_pointer, "weight")
|
290 |
+
encoder_pointer.weight = decoder_pointer.weight
|
291 |
+
if hasattr(decoder_pointer, "bias"):
|
292 |
+
assert hasattr(encoder_pointer, "bias")
|
293 |
+
encoder_pointer.bias = decoder_pointer.bias
|
294 |
+
print(module_name+' is tied')
|
295 |
+
return
|
296 |
+
|
297 |
+
encoder_modules = encoder_pointer._modules
|
298 |
+
decoder_modules = decoder_pointer._modules
|
299 |
+
if len(decoder_modules) > 0:
|
300 |
+
assert (
|
301 |
+
len(encoder_modules) > 0
|
302 |
+
), f"Encoder module {encoder_pointer} does not match decoder module {decoder_pointer}"
|
303 |
+
|
304 |
+
all_encoder_weights = set([module_name + "/" + sub_name for sub_name in encoder_modules.keys()])
|
305 |
+
encoder_layer_pos = 0
|
306 |
+
for name, module in decoder_modules.items():
|
307 |
+
if name.isdigit():
|
308 |
+
encoder_name = str(int(name) + encoder_layer_pos)
|
309 |
+
decoder_name = name
|
310 |
+
if not isinstance(decoder_modules[decoder_name], type(encoder_modules[encoder_name])) and len(
|
311 |
+
encoder_modules
|
312 |
+
) != len(decoder_modules):
|
313 |
+
# this can happen if the name corresponds to the position in a list module list of layers
|
314 |
+
# in this case the decoder has added a cross-attention that the encoder does not have
|
315 |
+
# thus skip this step and subtract one layer pos from encoder
|
316 |
+
encoder_layer_pos -= 1
|
317 |
+
continue
|
318 |
+
elif name not in encoder_modules:
|
319 |
+
continue
|
320 |
+
elif depth > 500:
|
321 |
+
raise ValueError(
|
322 |
+
"Max depth of recursive function `tie_encoder_to_decoder` reached. It seems that there is a circular dependency between two or more `nn.Modules` of your model."
|
323 |
+
)
|
324 |
+
else:
|
325 |
+
decoder_name = encoder_name = name
|
326 |
+
tie_encoder_to_decoder_recursively(
|
327 |
+
decoder_modules[decoder_name],
|
328 |
+
encoder_modules[encoder_name],
|
329 |
+
module_name + "/" + name,
|
330 |
+
uninitialized_encoder_weights,
|
331 |
+
skip_key,
|
332 |
+
depth=depth + 1,
|
333 |
+
)
|
334 |
+
all_encoder_weights.remove(module_name + "/" + encoder_name)
|
335 |
+
|
336 |
+
uninitialized_encoder_weights += list(all_encoder_weights)
|
337 |
+
|
338 |
+
# tie weights recursively
|
339 |
+
tie_encoder_to_decoder_recursively(decoder, encoder, base_model_prefix, uninitialized_encoder_weights, skip_key)
|