DoruC/Grounded-Segment-Anything
0
1# coding=utf-82# Copyright 2021 The UCLA NLP Authors and The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15""" PyTorch VisualBERT model."""16 17 18import math19from dataclasses import dataclass20from typing import Optional, Tuple, Union21 22import torch23import torch.utils.checkpoint24from torch import nn25from torch.nn import CrossEntropyLoss, KLDivLoss, LogSoftmax26 27from ...activations import ACT2FN28from ...modeling_outputs import (29 BaseModelOutput,30 BaseModelOutputWithPooling,31 MultipleChoiceModelOutput,32 SequenceClassifierOutput,33)34from ...modeling_utils import PreTrainedModel35from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer36from ...utils import (37 ModelOutput,38 add_start_docstrings,39 add_start_docstrings_to_model_forward,40 logging,41 replace_return_docstrings,42)43from .configuration_visual_bert import VisualBertConfig44 45 46logger = logging.get_logger(__name__)47 48_CONFIG_FOR_DOC = "VisualBertConfig"49_CHECKPOINT_FOR_DOC = "uclanlp/visualbert-vqa-coco-pre"50 51VISUAL_BERT_PRETRAINED_MODEL_ARCHIVE_LIST = [52 "uclanlp/visualbert-vqa",53 "uclanlp/visualbert-vqa-pre",54 "uclanlp/visualbert-vqa-coco-pre",55 "uclanlp/visualbert-vcr",56 "uclanlp/visualbert-vcr-pre",57 "uclanlp/visualbert-vcr-coco-pre",58 "uclanlp/visualbert-nlvr2",59 "uclanlp/visualbert-nlvr2-pre",60 "uclanlp/visualbert-nlvr2-coco-pre"61 # See all VisualBERT models at https://huggingface.co/models?filter=visual_bert62]63 64 65class VisualBertEmbeddings(nn.Module):66 """Construct the embeddings from word, position and token_type embeddings and visual embeddings."""67 68 def __init__(self, config):69 super().__init__()70 self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)71 self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)72 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)73 74 # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load75 # any TensorFlow checkpoint file76 77 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)78 self.dropout = nn.Dropout(config.hidden_dropout_prob)79 80 # position_ids (1, len position emb) is contiguous in memory and exported when serialized81 self.register_buffer(82 "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False83 )84 85 # For Visual Features86 # Token type and position embedding for image features87 self.visual_token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)88 self.visual_position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)89 90 if config.special_visual_initialize:91 self.visual_token_type_embeddings.weight.data = nn.Parameter(92 self.token_type_embeddings.weight.data.clone(), requires_grad=True93 )94 self.visual_position_embeddings.weight.data = nn.Parameter(95 self.position_embeddings.weight.data.clone(), requires_grad=True96 )97 98 self.visual_projection = nn.Linear(config.visual_embedding_dim, config.hidden_size)99 100 def forward(101 self,102 input_ids=None,103 token_type_ids=None,104 position_ids=None,105 inputs_embeds=None,106 visual_embeds=None,107 visual_token_type_ids=None,108 image_text_alignment=None,109 ):110 if input_ids is not None:111 input_shape = input_ids.size()112 else:113 input_shape = inputs_embeds.size()[:-1]114 115 seq_length = input_shape[1]116 117 if position_ids is None:118 position_ids = self.position_ids[:, :seq_length]119 120 if inputs_embeds is None:121 inputs_embeds = self.word_embeddings(input_ids)122 123 if token_type_ids is None:124 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)125 126 token_type_embeddings = self.token_type_embeddings(token_type_ids)127 128 embeddings = inputs_embeds + token_type_embeddings129 130 # Absolute Position Embeddings131 position_embeddings = self.position_embeddings(position_ids)132 embeddings += position_embeddings133 134 if visual_embeds is not None:135 if visual_token_type_ids is None:136 visual_token_type_ids = torch.ones(137 visual_embeds.size()[:-1], dtype=torch.long, device=self.position_ids.device138 )139 140 visual_embeds = self.visual_projection(visual_embeds)141 visual_token_type_embeddings = self.visual_token_type_embeddings(visual_token_type_ids)142 143 if image_text_alignment is not None:144 # image_text_alignment = Batch x image_length x alignment_number.145 # Each element denotes the position of the word corresponding to the image feature. -1 is the padding value.146 147 dtype = token_type_embeddings.dtype148 image_text_alignment_mask = (image_text_alignment != -1).long()149 # Get rid of the -1.150 image_text_alignment = image_text_alignment_mask * image_text_alignment151 152 # Batch x image_length x alignment length x dim153 visual_position_embeddings = self.position_embeddings(image_text_alignment)154 visual_position_embeddings *= image_text_alignment_mask.to(dtype=dtype).unsqueeze(-1)155 visual_position_embeddings = visual_position_embeddings.sum(2)156 157 # We want to averge along the alignment_number dimension.158 image_text_alignment_mask = image_text_alignment_mask.to(dtype=dtype).sum(2)159 160 if (image_text_alignment_mask == 0).sum() != 0:161 image_text_alignment_mask[image_text_alignment_mask == 0] = 1 # Avoid divide by zero error162 logger.warning(163 "Found 0 values in `image_text_alignment_mask`. Setting them to 1 to avoid divide-by-zero"164 " error."165 )166 visual_position_embeddings = visual_position_embeddings / image_text_alignment_mask.unsqueeze(-1)167 168 visual_position_ids = torch.zeros(169 *visual_embeds.size()[:-1], dtype=torch.long, device=visual_embeds.device170 )171 172 # When fine-tuning the detector , the image_text_alignment is sometimes padded too long.173 if visual_position_embeddings.size(1) != visual_embeds.size(1):174 if visual_position_embeddings.size(1) < visual_embeds.size(1):175 raise ValueError(176 f"Visual position embeddings length: {visual_position_embeddings.size(1)} "177 f"should be the same as `visual_embeds` length: {visual_embeds.size(1)}"178 )179 visual_position_embeddings = visual_position_embeddings[:, : visual_embeds.size(1), :]180 181 visual_position_embeddings = visual_position_embeddings + self.visual_position_embeddings(182 visual_position_ids183 )184 else:185 visual_position_ids = torch.zeros(186 *visual_embeds.size()[:-1], dtype=torch.long, device=visual_embeds.device187 )188 visual_position_embeddings = self.visual_position_embeddings(visual_position_ids)189 190 visual_embeddings = visual_embeds + visual_position_embeddings + visual_token_type_embeddings191 192 embeddings = torch.cat((embeddings, visual_embeddings), dim=1)193 194 embeddings = self.LayerNorm(embeddings)195 embeddings = self.dropout(embeddings)196 return embeddings197 198 199class VisualBertSelfAttention(nn.Module):200 def __init__(self, config):201 super().__init__()202 if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):203 raise ValueError(204 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "205 f"heads ({config.num_attention_heads})"206 )207 208 self.num_attention_heads = config.num_attention_heads209 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)210 self.all_head_size = self.num_attention_heads * self.attention_head_size211 212 self.query = nn.Linear(config.hidden_size, self.all_head_size)213 self.key = nn.Linear(config.hidden_size, self.all_head_size)214 self.value = nn.Linear(config.hidden_size, self.all_head_size)215 216 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)217 218 def transpose_for_scores(self, x):219 new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)220 x = x.view(*new_x_shape)221 return x.permute(0, 2, 1, 3)222 223 def forward(224 self,225 hidden_states,226 attention_mask=None,227 head_mask=None,228 output_attentions=False,229 ):230 mixed_query_layer = self.query(hidden_states)231 232 key_layer = self.transpose_for_scores(self.key(hidden_states))233 value_layer = self.transpose_for_scores(self.value(hidden_states))234 235 query_layer = self.transpose_for_scores(mixed_query_layer)236 237 # Take the dot product between "query" and "key" to get the raw attention scores.238 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))239 240 attention_scores = attention_scores / math.sqrt(self.attention_head_size)241 if attention_mask is not None:242 # Apply the attention mask is (precomputed for all layers in VisualBertSelfAttentionModel forward() function)243 attention_scores = attention_scores + attention_mask244 245 # Normalize the attention scores to probabilities.246 attention_probs = nn.functional.softmax(attention_scores, dim=-1)247 248 # This is actually dropping out entire tokens to attend to, which might249 # seem a bit unusual, but is taken from the original Transformer paper.250 attention_probs = self.dropout(attention_probs)251 252 # Mask heads if we want to253 if head_mask is not None:254 attention_probs = attention_probs * head_mask255 256 context_layer = torch.matmul(attention_probs, value_layer)257 258 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()259 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)260 context_layer = context_layer.view(*new_context_layer_shape)261 262 outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)263 264 return outputs265 266 267# Copied from transformers.models.bert.modeling_bert.BertSelfOutput with Bert->VisualBert268class VisualBertSelfOutput(nn.Module):269 def __init__(self, config):270 super().__init__()271 self.dense = nn.Linear(config.hidden_size, config.hidden_size)272 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)273 self.dropout = nn.Dropout(config.hidden_dropout_prob)274 275 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:276 hidden_states = self.dense(hidden_states)277 hidden_states = self.dropout(hidden_states)278 hidden_states = self.LayerNorm(hidden_states + input_tensor)279 return hidden_states280 281 282class VisualBertAttention(nn.Module):283 def __init__(self, config):284 super().__init__()285 self.self = VisualBertSelfAttention(config)286 self.output = VisualBertSelfOutput(config)287 self.pruned_heads = set()288 289 def prune_heads(self, heads):290 if len(heads) == 0:291 return292 heads, index = find_pruneable_heads_and_indices(293 heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads294 )295 296 # Prune linear layers297 self.self.query = prune_linear_layer(self.self.query, index)298 self.self.key = prune_linear_layer(self.self.key, index)299 self.self.value = prune_linear_layer(self.self.value, index)300 self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)301 302 # Update hyper params and store pruned heads303 self.self.num_attention_heads = self.self.num_attention_heads - len(heads)304 self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads305 self.pruned_heads = self.pruned_heads.union(heads)306 307 def forward(308 self,309 hidden_states,310 attention_mask=None,311 head_mask=None,312 output_attentions=False,313 ):314 self_outputs = self.self(315 hidden_states,316 attention_mask,317 head_mask,318 output_attentions,319 )320 attention_output = self.output(self_outputs[0], hidden_states)321 outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them322 return outputs323 324 325# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->VisualBert326class VisualBertIntermediate(nn.Module):327 def __init__(self, config):328 super().__init__()329 self.dense = nn.Linear(config.hidden_size, config.intermediate_size)330 if isinstance(config.hidden_act, str):331 self.intermediate_act_fn = ACT2FN[config.hidden_act]332 else:333 self.intermediate_act_fn = config.hidden_act334 335 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:336 hidden_states = self.dense(hidden_states)337 hidden_states = self.intermediate_act_fn(hidden_states)338 return hidden_states339 340 341# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->VisualBert342class VisualBertOutput(nn.Module):343 def __init__(self, config):344 super().__init__()345 self.dense = nn.Linear(config.intermediate_size, config.hidden_size)346 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)347 self.dropout = nn.Dropout(config.hidden_dropout_prob)348 349 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:350 hidden_states = self.dense(hidden_states)351 hidden_states = self.dropout(hidden_states)352 hidden_states = self.LayerNorm(hidden_states + input_tensor)353 return hidden_states354 355 356class VisualBertLayer(nn.Module):357 def __init__(self, config):358 super().__init__()359 self.chunk_size_feed_forward = config.chunk_size_feed_forward360 self.seq_len_dim = 1361 self.attention = VisualBertAttention(config)362 self.intermediate = VisualBertIntermediate(config)363 self.output = VisualBertOutput(config)364 365 def forward(366 self,367 hidden_states,368 attention_mask=None,369 head_mask=None,370 output_attentions=False,371 ):372 self_attention_outputs = self.attention(373 hidden_states,374 attention_mask,375 head_mask,376 output_attentions=output_attentions,377 )378 attention_output = self_attention_outputs[0]379 380 outputs = self_attention_outputs[1:] # add self attentions if we output attention weights381 382 layer_output = apply_chunking_to_forward(383 self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output384 )385 outputs = (layer_output,) + outputs386 387 return outputs388 389 def feed_forward_chunk(self, attention_output):390 intermediate_output = self.intermediate(attention_output)391 layer_output = self.output(intermediate_output, attention_output)392 return layer_output393 394 395class VisualBertEncoder(nn.Module):396 def __init__(self, config):397 super().__init__()398 self.config = config399 self.layer = nn.ModuleList([VisualBertLayer(config) for _ in range(config.num_hidden_layers)])400 self.gradient_checkpointing = False401 402 def forward(403 self,404 hidden_states,405 attention_mask=None,406 head_mask=None,407 output_attentions=False,408 output_hidden_states=False,409 return_dict=True,410 ):411 all_hidden_states = () if output_hidden_states else None412 all_self_attentions = () if output_attentions else None413 414 for i, layer_module in enumerate(self.layer):415 if output_hidden_states:416 all_hidden_states = all_hidden_states + (hidden_states,)417 418 layer_head_mask = head_mask[i] if head_mask is not None else None419 420 if self.gradient_checkpointing and self.training:421 422 def create_custom_forward(module):423 def custom_forward(*inputs):424 return module(*inputs, output_attentions)425 426 return custom_forward427 428 layer_outputs = torch.utils.checkpoint.checkpoint(429 create_custom_forward(layer_module),430 hidden_states,431 attention_mask,432 layer_head_mask,433 )434 else:435 layer_outputs = layer_module(hidden_states, attention_mask, layer_head_mask, output_attentions)436 437 hidden_states = layer_outputs[0]438 if output_attentions:439 all_self_attentions = all_self_attentions + (layer_outputs[1],)440 441 if output_hidden_states:442 all_hidden_states = all_hidden_states + (hidden_states,)443 444 if not return_dict:445 return tuple(446 v447 for v in [448 hidden_states,449 all_hidden_states,450 all_self_attentions,451 ]452 if v is not None453 )454 return BaseModelOutput(455 last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_self_attentions456 )457 458 459# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->VisualBert460class VisualBertPooler(nn.Module):461 def __init__(self, config):462 super().__init__()463 self.dense = nn.Linear(config.hidden_size, config.hidden_size)464 self.activation = nn.Tanh()465 466 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:467 # We "pool" the model by simply taking the hidden state corresponding468 # to the first token.469 first_token_tensor = hidden_states[:, 0]470 pooled_output = self.dense(first_token_tensor)471 pooled_output = self.activation(pooled_output)472 return pooled_output473 474 475# Copied from transformers.models.bert.modeling_bert.BertPredictionHeadTransform with Bert->VisualBert476class VisualBertPredictionHeadTransform(nn.Module):477 def __init__(self, config):478 super().__init__()479 self.dense = nn.Linear(config.hidden_size, config.hidden_size)480 if isinstance(config.hidden_act, str):481 self.transform_act_fn = ACT2FN[config.hidden_act]482 else:483 self.transform_act_fn = config.hidden_act484 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)485 486 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:487 hidden_states = self.dense(hidden_states)488 hidden_states = self.transform_act_fn(hidden_states)489 hidden_states = self.LayerNorm(hidden_states)490 return hidden_states491 492 493# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead with Bert->VisualBert494class VisualBertLMPredictionHead(nn.Module):495 def __init__(self, config):496 super().__init__()497 self.transform = VisualBertPredictionHeadTransform(config)498 499 # The output weights are the same as the input embeddings, but there is500 # an output-only bias for each token.501 self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)502 503 self.bias = nn.Parameter(torch.zeros(config.vocab_size))504 505 # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`506 self.decoder.bias = self.bias507 508 def forward(self, hidden_states):509 hidden_states = self.transform(hidden_states)510 hidden_states = self.decoder(hidden_states)511 return hidden_states512 513 514# Copied from transformers.models.bert.modeling_bert.BertPreTrainingHeads with Bert->VisualBert515class VisualBertPreTrainingHeads(nn.Module):516 def __init__(self, config):517 super().__init__()518 self.predictions = VisualBertLMPredictionHead(config)519 self.seq_relationship = nn.Linear(config.hidden_size, 2)520 521 def forward(self, sequence_output, pooled_output):522 prediction_scores = self.predictions(sequence_output)523 seq_relationship_score = self.seq_relationship(pooled_output)524 return prediction_scores, seq_relationship_score525 526 527class VisualBertPreTrainedModel(PreTrainedModel):528 """529 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained530 models.531 """532 533 config_class = VisualBertConfig534 base_model_prefix = "visual_bert"535 supports_gradient_checkpointing = True536 537 def _init_weights(self, module):538 """Initialize the weights"""539 if isinstance(module, (nn.Linear, nn.Embedding)):540 # Slightly different from the TF version which uses truncated_normal for initialization541 # cf https://github.com/pytorch/pytorch/pull/5617542 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)543 544 elif isinstance(module, nn.LayerNorm):545 module.bias.data.zero_()546 module.weight.data.fill_(1.0)547 if isinstance(module, nn.Linear) and module.bias is not None:548 module.bias.data.zero_()549 550 def _set_gradient_checkpointing(self, module, value=False):551 if isinstance(module, VisualBertEncoder):552 module.gradient_checkpointing = value553 554 555@dataclass556class VisualBertForPreTrainingOutput(ModelOutput):557 """558 Output type of [`VisualBertForPreTraining`].559 560 Args:561 loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):562 Total loss as the sum of the masked language modeling loss and the sentence-image prediction563 (classification) loss.564 prediction_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):565 Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).566 seq_relationship_logits (`torch.FloatTensor` of shape `(batch_size, 2)`):567 Prediction scores of the sentence-image prediction (classification) head (scores of True/False continuation568 before SoftMax).569 hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):570 Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of571 shape `(batch_size, sequence_length, hidden_size)`.572 573 Hidden-states of the model at the output of each layer plus the initial embedding outputs.574 attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):575 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,576 sequence_length)`.577 578 Attentions weights after the attention softmax, used to compute the weighted average in the self-attention579 heads.580 """581 582 loss: Optional[torch.FloatTensor] = None583 prediction_logits: torch.FloatTensor = None584 seq_relationship_logits: torch.FloatTensor = None585 hidden_states: Optional[Tuple[torch.FloatTensor]] = None586 attentions: Optional[Tuple[torch.FloatTensor]] = None587 588 589VISUAL_BERT_START_DOCSTRING = r"""590 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the591 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads592 etc.)593 594 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.595 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage596 and behavior.597 598 Parameters:599 config ([`VisualBertConfig`]): Model configuration class with all the parameters of the model.600 Initializing with a config file does not load the weights associated with the model, only the601 configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.602"""603 604VISUAL_BERT_INPUTS_DOCSTRING = r"""605 Args:606 input_ids (`torch.LongTensor` of shape `({0})`):607 Indices of input sequence tokens in the vocabulary.608 609 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and610 [`PreTrainedTokenizer.__call__`] for details.611 612 [What are input IDs?](../glossary#input-ids)613 attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):614 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:615 616 - 1 for tokens that are **not masked**,617 - 0 for tokens that are **masked**.618 619 [What are attention masks?](../glossary#attention-mask)620 token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):621 Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,622 1]`:623 624 - 0 corresponds to a *sentence A* token,625 - 1 corresponds to a *sentence B* token.626 627 [What are token type IDs?](../glossary#token-type-ids)628 position_ids (`torch.LongTensor` of shape `({0})`, *optional*):629 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,630 config.max_position_embeddings - 1]`.631 632 [What are position IDs?](../glossary#position-ids)633 head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):634 Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:635 636 - 1 indicates the head is **not masked**,637 - 0 indicates the head is **masked**.638 639 inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):640 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This641 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the642 model's internal embedding lookup matrix.643 644 visual_embeds (`torch.FloatTensor` of shape `(batch_size, visual_seq_length, visual_embedding_dim)`, *optional*):645 The embedded representation of the visual inputs, generally derived using using an object detector.646 647 visual_attention_mask (`torch.FloatTensor` of shape `(batch_size, visual_seq_length)`, *optional*):648 Mask to avoid performing attention on visual embeddings. Mask values selected in `[0, 1]`:649 650 - 1 for tokens that are **not masked**,651 - 0 for tokens that are **masked**.652 653 [What are attention masks?](../glossary#attention-mask)654 visual_token_type_ids (`torch.LongTensor` of shape `(batch_size, visual_seq_length)`, *optional*):655 Segment token indices to indicate different portions of the visual embeds.656 657 [What are token type IDs?](../glossary#token-type-ids) The authors of VisualBERT set the658 *visual_token_type_ids* to *1* for all tokens.659 660 image_text_alignment (`torch.LongTensor` of shape `(batch_size, visual_seq_length, alignment_number)`, *optional*):661 Image-Text alignment uses to decide the position IDs of the visual embeddings.662 663 output_attentions (`bool`, *optional*):664 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned665 tensors for more detail.666 output_hidden_states (`bool`, *optional*):667 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for668 more detail.669 return_dict (`bool`, *optional*):670 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.671"""672 673 674@add_start_docstrings(675 "The bare VisualBert Model transformer outputting raw hidden-states without any specific head on top.",676 VISUAL_BERT_START_DOCSTRING,677)678class VisualBertModel(VisualBertPreTrainedModel):679 """680 681 The model can behave as an encoder (with only self-attention) following the architecture described in [Attention is682 all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,683 Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.684 """685 686 def __init__(self, config, add_pooling_layer=True):687 super().__init__(config)688 self.config = config689 690 self.embeddings = VisualBertEmbeddings(config)691 self.encoder = VisualBertEncoder(config)692 693 self.pooler = VisualBertPooler(config) if add_pooling_layer else None694 695 self.bypass_transformer = config.bypass_transformer696 697 if self.bypass_transformer:698 self.additional_layer = VisualBertLayer(config)699 700 # Initialize weights and apply final processing701 self.post_init()702 703 def get_input_embeddings(self):704 return self.embeddings.word_embeddings705 706 def set_input_embeddings(self, value):707 self.embeddings.word_embeddings = value708 709 def _prune_heads(self, heads_to_prune):710 """711 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base712 class PreTrainedModel713 """714 for layer, heads in heads_to_prune.items():715 self.encoder.layer[layer].attention.prune_heads(heads)716 717 @add_start_docstrings_to_model_forward(VISUAL_BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))718 @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=_CONFIG_FOR_DOC)719 def forward(720 self,721 input_ids: Optional[torch.LongTensor] = None,722 attention_mask: Optional[torch.LongTensor] = None,723 token_type_ids: Optional[torch.LongTensor] = None,724 position_ids: Optional[torch.LongTensor] = None,725 head_mask: Optional[torch.LongTensor] = None,726 inputs_embeds: Optional[torch.FloatTensor] = None,727 visual_embeds: Optional[torch.FloatTensor] = None,728 visual_attention_mask: Optional[torch.LongTensor] = None,729 visual_token_type_ids: Optional[torch.LongTensor] = None,730 image_text_alignment: Optional[torch.LongTensor] = None,731 output_attentions: Optional[bool] = None,732 output_hidden_states: Optional[bool] = None,733 return_dict: Optional[bool] = None,734 ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPooling]:735 r"""736 737 Returns:738 739 Example:740 741 ```python742 # Assumption: *get_visual_embeddings(image)* gets the visual embeddings of the image.743 from transformers import AutoTokenizer, VisualBertModel744 import torch745 746 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")747 model = VisualBertModel.from_pretrained("uclanlp/visualbert-vqa-coco-pre")748 749 inputs = tokenizer("The capital of France is Paris.", return_tensors="pt")750 visual_embeds = get_visual_embeddings(image).unsqueeze(0)751 visual_token_type_ids = torch.ones(visual_embeds.shape[:-1], dtype=torch.long)752 visual_attention_mask = torch.ones(visual_embeds.shape[:-1], dtype=torch.float)753 754 inputs.update(755 {756 "visual_embeds": visual_embeds,757 "visual_token_type_ids": visual_token_type_ids,758 "visual_attention_mask": visual_attention_mask,759 }760 )761 762 outputs = model(**inputs)763 764 last_hidden_states = outputs.last_hidden_state765 ```"""766 767 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions768 output_hidden_states = (769 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states770 )771 return_dict = return_dict if return_dict is not None else self.config.use_return_dict772 773 if input_ids is not None and inputs_embeds is not None:774 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")775 elif input_ids is not None:776 self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)777 input_shape = input_ids.size()778 elif inputs_embeds is not None:779 input_shape = inputs_embeds.size()[:-1]780 else:781 raise ValueError("You have to specify either input_ids or inputs_embeds")782 783 batch_size, seq_length = input_shape784 device = input_ids.device if input_ids is not None else inputs_embeds.device785 786 if visual_embeds is not None:787 visual_input_shape = visual_embeds.size()[:-1]788 789 if attention_mask is None:790 attention_mask = torch.ones(input_shape, device=device)791 792 if visual_embeds is not None and visual_attention_mask is None:793 visual_attention_mask = torch.ones(visual_input_shape, device=device)794 795 # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]796 # ourselves in which case we just need to make it broadcastable to all heads.797 if visual_embeds is not None:798 combined_attention_mask = torch.cat((attention_mask, visual_attention_mask), dim=-1)799 extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(800 combined_attention_mask, (batch_size, input_shape + visual_input_shape)801 )802 803 else:804 extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(805 attention_mask, (batch_size, input_shape)806 )807 808 # Prepare head mask if needed809 # 1.0 in head_mask indicate we keep the head810 # attention_probs has shape bsz x n_heads x N x N811 # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]812 # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]813 head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)814 815 embedding_output = self.embeddings(816 input_ids=input_ids,817 position_ids=position_ids,818 token_type_ids=token_type_ids,819 inputs_embeds=inputs_embeds,820 visual_embeds=visual_embeds,821 visual_token_type_ids=visual_token_type_ids,822 image_text_alignment=image_text_alignment,823 )824 825 if self.bypass_transformer and visual_embeds is not None:826 text_length = input_ids.size(1)827 text_embedding_output = embedding_output[:, :text_length, :]828 visual_embedding_output = embedding_output[:, text_length:, :]829 830 text_extended_attention_mask = extended_attention_mask[:, :, text_length, :text_length]831 832 encoded_outputs = self.encoder(833 text_embedding_output,834 attention_mask=text_extended_attention_mask,835 output_attentions=output_attentions,836 output_hidden_states=output_hidden_states,837 return_dict=return_dict,838 )839 sequence_output = encoded_outputs[0]840 concatenated_input = torch.cat((sequence_output, visual_embedding_output), dim=1)841 sequence_output = self.additional_layer(concatenated_input, extended_attention_mask)842 pooled_output = self.pooler(sequence_output) if self.pooler is not None else None843 844 else:845 encoder_outputs = self.encoder(846 embedding_output,847 attention_mask=extended_attention_mask,848 head_mask=head_mask,849 output_attentions=output_attentions,850 output_hidden_states=output_hidden_states,851 return_dict=return_dict,852 )853 sequence_output = encoder_outputs[0]854 855 pooled_output = self.pooler(sequence_output) if self.pooler is not None else None856 857 if not return_dict:858 return (sequence_output, pooled_output) + encoder_outputs[1:]859 860 return BaseModelOutputWithPooling(861 last_hidden_state=sequence_output,862 pooler_output=pooled_output,863 hidden_states=encoder_outputs.hidden_states,864 attentions=encoder_outputs.attentions,865 )866 867 868@add_start_docstrings(869 """870 VisualBert Model with two heads on top as done during the pretraining: a `masked language modeling` head and a871 `sentence-image prediction (classification)` head.872 """,873 VISUAL_BERT_START_DOCSTRING,874)875class VisualBertForPreTraining(VisualBertPreTrainedModel):876 _tied_weights_keys = ["cls.predictions.decoder.weight", "cls.predictions.decoder.bias"]877 878 def __init__(self, config):879 super().__init__(config)880 881 self.visual_bert = VisualBertModel(config)882 self.cls = VisualBertPreTrainingHeads(config)883 884 # Initialize weights and apply final processing885 self.post_init()886 887 def get_output_embeddings(self):888 return self.cls.predictions.decoder889 890 def set_output_embeddings(self, new_embeddings):891 self.cls.predictions.decoder = new_embeddings892 893 @add_start_docstrings_to_model_forward(VISUAL_BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))894 @replace_return_docstrings(output_type=VisualBertForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)895 def forward(896 self,897 input_ids: Optional[torch.LongTensor] = None,898 attention_mask: Optional[torch.LongTensor] = None,899 token_type_ids: Optional[torch.LongTensor] = None,900 position_ids: Optional[torch.LongTensor] = None,901 head_mask: Optional[torch.LongTensor] = None,902 inputs_embeds: Optional[torch.FloatTensor] = None,903 visual_embeds: Optional[torch.FloatTensor] = None,904 visual_attention_mask: Optional[torch.LongTensor] = None,905 visual_token_type_ids: Optional[torch.LongTensor] = None,906 image_text_alignment: Optional[torch.LongTensor] = None,907 output_attentions: Optional[bool] = None,908 output_hidden_states: Optional[bool] = None,909 return_dict: Optional[bool] = None,910 labels: Optional[torch.LongTensor] = None,911 sentence_image_labels: Optional[torch.LongTensor] = None,912 ) -> Union[Tuple[torch.Tensor], VisualBertForPreTrainingOutput]:913 r"""914 labels (`torch.LongTensor` of shape `(batch_size, total_sequence_length)`, *optional*):915 Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,916 config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the917 loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`918 sentence_image_labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):919 Labels for computing the sentence-image prediction (classification) loss. Input should be a sequence pair920 (see `input_ids` docstring) Indices should be in `[0, 1]`:921 922 - 0 indicates sequence B is a matching pair of sequence A for the given image,923 - 1 indicates sequence B is a random sequence w.r.t A for the given image.924 925 Returns:926 927 Example:928 929 ```python930 # Assumption: *get_visual_embeddings(image)* gets the visual embeddings of the image in the batch.931 from transformers import AutoTokenizer, VisualBertForPreTraining932 933 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")934 model = VisualBertForPreTraining.from_pretrained("uclanlp/visualbert-vqa-coco-pre")935 936 inputs = tokenizer("The capital of France is [MASK].", return_tensors="pt")937 visual_embeds = get_visual_embeddings(image).unsqueeze(0)938 visual_token_type_ids = torch.ones(visual_embeds.shape[:-1], dtype=torch.long)939 visual_attention_mask = torch.ones(visual_embeds.shape[:-1], dtype=torch.float)940 941 inputs.update(942 {943 "visual_embeds": visual_embeds,944 "visual_token_type_ids": visual_token_type_ids,945 "visual_attention_mask": visual_attention_mask,946 }947 )948 max_length = inputs["input_ids"].shape[-1] + visual_embeds.shape[-2]949 labels = tokenizer(950 "The capital of France is Paris.", return_tensors="pt", padding="max_length", max_length=max_length951 )["input_ids"]952 sentence_image_labels = torch.tensor(1).unsqueeze(0) # Batch_size953 954 955 outputs = model(**inputs, labels=labels, sentence_image_labels=sentence_image_labels)956 loss = outputs.loss957 prediction_logits = outputs.prediction_logits958 seq_relationship_logits = outputs.seq_relationship_logits959 ```"""960 return_dict = return_dict if return_dict is not None else self.config.use_return_dict961 962 outputs = self.visual_bert(963 input_ids,964 attention_mask=attention_mask,965 token_type_ids=token_type_ids,966 position_ids=position_ids,967 head_mask=head_mask,968 inputs_embeds=inputs_embeds,969 visual_embeds=visual_embeds,970 visual_attention_mask=visual_attention_mask,971 visual_token_type_ids=visual_token_type_ids,972 image_text_alignment=image_text_alignment,973 output_attentions=output_attentions,974 output_hidden_states=output_hidden_states,975 return_dict=return_dict,976 )977 978 sequence_output, pooled_output = outputs[:2]979 prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)980 981 total_loss = None982 if labels is not None and sentence_image_labels is not None:983 total_size = attention_mask.size(-1) + visual_attention_mask.size(-1)984 if labels.size(-1) != total_size:985 raise ValueError(986 "The labels provided should have same sequence length as total attention mask. "987 f"Found labels with sequence length {labels.size(-1)}, expected {total_size}."988 )989 990 loss_fct = CrossEntropyLoss()991 masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))992 sentence_image_loss = loss_fct(seq_relationship_score.view(-1, 2), sentence_image_labels.view(-1))993 total_loss = masked_lm_loss + sentence_image_loss994 995 if labels is not None and sentence_image_labels is None:996 total_size = attention_mask.size(-1) + visual_attention_mask.size(-1)997 if labels.size(-1) != total_size:998 raise ValueError(999 "The labels provided should have same sequence length as total attention mask. "1000 f"Found labels with sequence length {labels.size(-1)}, expected {total_size}."1001 )1002 1003 loss_fct = CrossEntropyLoss()1004 total_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))1005 1006 if not return_dict:1007 output = (prediction_scores, seq_relationship_score) + outputs[2:]1008 return ((total_loss,) + output) if total_loss is not None else output1009 1010 return VisualBertForPreTrainingOutput(1011 loss=total_loss,1012 prediction_logits=prediction_scores,1013 seq_relationship_logits=seq_relationship_score,1014 hidden_states=outputs.hidden_states,1015 attentions=outputs.attentions,1016 )1017 1018 1019@add_start_docstrings(1020 """1021 VisualBert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and1022 a softmax) e.g. for VCR tasks.1023 """,1024 VISUAL_BERT_START_DOCSTRING,1025)1026class VisualBertForMultipleChoice(VisualBertPreTrainedModel):1027 def __init__(self, config):1028 super().__init__(config)1029 1030 self.visual_bert = VisualBertModel(config)1031 self.dropout = nn.Dropout(config.hidden_dropout_prob)1032 self.cls = nn.Linear(config.hidden_size, 1)1033 1034 # Initialize weights and apply final processing1035 self.post_init()1036 1037 @add_start_docstrings_to_model_forward(1038 VISUAL_BERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length")1039 )1040 @replace_return_docstrings(output_type=MultipleChoiceModelOutput, config_class=_CONFIG_FOR_DOC)1041 def forward(1042 self,1043 input_ids: Optional[torch.LongTensor] = None,1044 attention_mask: Optional[torch.LongTensor] = None,1045 token_type_ids: Optional[torch.LongTensor] = None,1046 position_ids: Optional[torch.LongTensor] = None,1047 head_mask: Optional[torch.LongTensor] = None,1048 inputs_embeds: Optional[torch.FloatTensor] = None,1049 visual_embeds: Optional[torch.FloatTensor] = None,1050 visual_attention_mask: Optional[torch.LongTensor] = None,1051 visual_token_type_ids: Optional[torch.LongTensor] = None,1052 image_text_alignment: Optional[torch.LongTensor] = None,1053 output_attentions: Optional[bool] = None,1054 output_hidden_states: Optional[bool] = None,1055 return_dict: Optional[bool] = None,1056 labels: Optional[torch.LongTensor] = None,1057 ) -> Union[Tuple[torch.Tensor], MultipleChoiceModelOutput]:1058 r"""1059 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1060 Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,1061 num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See1062 `input_ids` above)1063 1064 Returns:1065 1066 Example:1067 1068 ```python1069 # Assumption: *get_visual_embeddings(image)* gets the visual embeddings of the image in the batch.1070 from transformers import AutoTokenizer, VisualBertForMultipleChoice1071 import torch1072 1073 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")1074 model = VisualBertForMultipleChoice.from_pretrained("uclanlp/visualbert-vcr")1075 1076 prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."1077 choice0 = "It is eaten with a fork and a knife."1078 choice1 = "It is eaten while held in the hand."1079 1080 visual_embeds = get_visual_embeddings(image)1081 # (batch_size, num_choices, visual_seq_length, visual_embedding_dim)1082 visual_embeds = visual_embeds.expand(1, 2, *visual_embeds.shape)1083 visual_token_type_ids = torch.ones(visual_embeds.shape[:-1], dtype=torch.long)1084 visual_attention_mask = torch.ones(visual_embeds.shape[:-1], dtype=torch.float)1085 1086 labels = torch.tensor(0).unsqueeze(0) # choice0 is correct (according to Wikipedia ;)), batch size 11087 1088 encoding = tokenizer([[prompt, prompt], [choice0, choice1]], return_tensors="pt", padding=True)1089 # batch size is 11090 inputs_dict = {k: v.unsqueeze(0) for k, v in encoding.items()}1091 inputs_dict.update(1092 {1093 "visual_embeds": visual_embeds,1094 "visual_attention_mask": visual_attention_mask,1095 "visual_token_type_ids": visual_token_type_ids,1096 "labels": labels,1097 }1098 )1099 outputs = model(**inputs_dict)1100 1101 loss = outputs.loss1102 logits = outputs.logits1103 ```"""1104 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1105 num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]1106 1107 input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None1108 attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None1109 token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None1110 position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None1111 inputs_embeds = (1112 inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))1113 if inputs_embeds is not None1114 else None1115 )1116 1117 visual_embeds = (1118 visual_embeds.view(-1, visual_embeds.size(-2), visual_embeds.size(-1))1119 if visual_embeds is not None1120 else None1121 )1122 visual_attention_mask = (1123 visual_attention_mask.view(-1, visual_attention_mask.size(-1))1124 if visual_attention_mask is not None1125 else None1126 )1127 visual_token_type_ids = (1128 visual_token_type_ids.view(-1, visual_token_type_ids.size(-1))1129 if visual_token_type_ids is not None1130 else None1131 )1132 1133 outputs = self.visual_bert(1134 input_ids,1135 attention_mask=attention_mask,1136 token_type_ids=token_type_ids,1137 position_ids=position_ids,1138 head_mask=head_mask,1139 inputs_embeds=inputs_embeds,1140 visual_embeds=visual_embeds,1141 visual_attention_mask=visual_attention_mask,1142 visual_token_type_ids=visual_token_type_ids,1143 image_text_alignment=image_text_alignment,1144 output_attentions=output_attentions,1145 output_hidden_states=output_hidden_states,1146 return_dict=return_dict,1147 )1148 1149 _, pooled_output = outputs[0], outputs[1]1150 1151 pooled_output = self.dropout(pooled_output)1152 logits = self.cls(pooled_output)1153 reshaped_logits = logits.view(-1, num_choices)1154 1155 loss = None1156 if labels is not None:1157 loss_fct = CrossEntropyLoss()1158 loss = loss_fct(reshaped_logits, labels)1159 1160 if not return_dict:1161 output = (reshaped_logits,) + outputs[2:]1162 return ((loss,) + output) if loss is not None else output1163 1164 return MultipleChoiceModelOutput(1165 loss=loss,1166 logits=reshaped_logits,1167 hidden_states=outputs.hidden_states,1168 attentions=outputs.attentions,1169 )1170 1171 1172@add_start_docstrings(1173 """1174 VisualBert Model with a classification/regression head on top (a dropout and a linear layer on top of the pooled1175 output) for VQA.1176 """,1177 VISUAL_BERT_START_DOCSTRING,1178)1179class VisualBertForQuestionAnswering(VisualBertPreTrainedModel):1180 def __init__(self, config):1181 super().__init__(config)1182 self.num_labels = config.num_labels1183 1184 self.visual_bert = VisualBertModel(config)1185 self.dropout = nn.Dropout(config.hidden_dropout_prob)1186 self.cls = nn.Linear(config.hidden_size, config.num_labels)1187 1188 # Initialize weights and apply final processing1189 self.post_init()1190 1191 @add_start_docstrings_to_model_forward(VISUAL_BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))1192 @replace_return_docstrings(output_type=SequenceClassifierOutput, config_class=_CONFIG_FOR_DOC)1193 def forward(1194 self,1195 input_ids: Optional[torch.LongTensor] = None,1196 attention_mask: Optional[torch.LongTensor] = None,1197 token_type_ids: Optional[torch.LongTensor] = None,1198 position_ids: Optional[torch.LongTensor] = None,1199 head_mask: Optional[torch.LongTensor] = None,1200 inputs_embeds: Optional[torch.FloatTensor] = None,