doevent commited on
Commit
eb5e503
1 Parent(s): d1e5ec4

Upload models/blip_pretrain.py

Browse files
Files changed (1) hide show
  1. 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)