Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2018 Hao Tan, Mohit Bansal, and the HuggingFace team3#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 LXMERT model."""16 17import math18import os19import warnings20from dataclasses import dataclass21from typing import Optional, Union22 23import torch24from torch import nn25from torch.nn import CrossEntropyLoss, SmoothL1Loss26 27from ...activations import ACT2FN, gelu28from ...modeling_utils import PreTrainedModel29from ...utils import ModelOutput, auto_docstring, logging30from .configuration_lxmert import LxmertConfig31 32 33logger = logging.get_logger(__name__)34 35 36class GeLU(nn.Module):37 def __init__(self):38 super().__init__()39 40 def forward(self, x):41 return gelu(x)42 43 44@dataclass45@auto_docstring(46 custom_intro="""47 Lxmert's outputs that contain the last hidden states, pooled outputs, and attention probabilities for the language,48 visual, and, cross-modality encoders. (note: the visual encoder in Lxmert is referred to as the "relation-ship"49 encoder")50 """51)52class LxmertModelOutput(ModelOutput):53 r"""54 language_output (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):55 Sequence of hidden-states at the output of the last layer of the language encoder.56 vision_output (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):57 Sequence of hidden-states at the output of the last layer of the visual encoder.58 pooled_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`):59 Last layer hidden-state of the first token of the sequence (classification, CLS, token) further processed60 by a Linear layer and a Tanh activation function. The Linear61 language_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):62 Tuple of `torch.FloatTensor` (one for input features + one for the output of each cross-modality layer) of63 shape `(batch_size, sequence_length, hidden_size)`.64 vision_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):65 Tuple of `torch.FloatTensor` (one for input features + one for the output of each cross-modality layer) of66 shape `(batch_size, sequence_length, hidden_size)`.67 language_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):68 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,69 sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in70 the self-attention heads.71 vision_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):72 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,73 sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in74 the self-attention heads.75 cross_encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):76 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,77 sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in78 the self-attention heads.79 """80 81 language_output: Optional[torch.FloatTensor] = None82 vision_output: Optional[torch.FloatTensor] = None83 pooled_output: Optional[torch.FloatTensor] = None84 language_hidden_states: Optional[tuple[torch.FloatTensor]] = None85 vision_hidden_states: Optional[tuple[torch.FloatTensor]] = None86 language_attentions: Optional[tuple[torch.FloatTensor]] = None87 vision_attentions: Optional[tuple[torch.FloatTensor]] = None88 cross_encoder_attentions: Optional[tuple[torch.FloatTensor]] = None89 90 91@dataclass92@auto_docstring(93 custom_intro="""94 Output type of [`LxmertForQuestionAnswering`].95 """96)97class LxmertForQuestionAnsweringOutput(ModelOutput):98 r"""99 loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):100 Total loss as the sum of the masked language modeling loss and the next sequence prediction101 (classification) loss.k.102 question_answering_score (`torch.FloatTensor` of shape `(batch_size, n_qa_answers)`, *optional*):103 Prediction scores of question answering objective (classification).104 language_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):105 Tuple of `torch.FloatTensor` (one for input features + one for the output of each cross-modality layer) of106 shape `(batch_size, sequence_length, hidden_size)`.107 vision_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):108 Tuple of `torch.FloatTensor` (one for input features + one for the output of each cross-modality layer) of109 shape `(batch_size, sequence_length, hidden_size)`.110 language_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):111 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,112 sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in113 the self-attention heads.114 vision_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):115 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,116 sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in117 the self-attention heads.118 cross_encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):119 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,120 sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in121 the self-attention heads.122 """123 124 loss: Optional[torch.FloatTensor] = None125 question_answering_score: Optional[torch.FloatTensor] = None126 language_hidden_states: Optional[tuple[torch.FloatTensor]] = None127 vision_hidden_states: Optional[tuple[torch.FloatTensor]] = None128 language_attentions: Optional[tuple[torch.FloatTensor]] = None129 vision_attentions: Optional[tuple[torch.FloatTensor]] = None130 cross_encoder_attentions: Optional[tuple[torch.FloatTensor]] = None131 132 133@dataclass134@auto_docstring(135 custom_intro="""136 Output type of [`LxmertForPreTraining`].137 """138)139class LxmertForPreTrainingOutput(ModelOutput):140 r"""141 loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):142 Total loss as the sum of the masked language modeling loss and the next sequence prediction143 (classification) loss.144 prediction_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):145 Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).146 cross_relationship_score (`torch.FloatTensor` of shape `(batch_size, 2)`):147 Prediction scores of the textual matching objective (classification) head (scores of True/False148 continuation before SoftMax).149 question_answering_score (`torch.FloatTensor` of shape `(batch_size, n_qa_answers)`):150 Prediction scores of question answering objective (classification).151 language_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):152 Tuple of `torch.FloatTensor` (one for input features + one for the output of each cross-modality layer) of153 shape `(batch_size, sequence_length, hidden_size)`.154 vision_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):155 Tuple of `torch.FloatTensor` (one for input features + one for the output of each cross-modality layer) of156 shape `(batch_size, sequence_length, hidden_size)`.157 language_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):158 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,159 sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in160 the self-attention heads.161 vision_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):162 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,163 sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in164 the self-attention heads.165 cross_encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):166 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,167 sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in168 the self-attention heads.169 """170 171 loss: Optional[torch.FloatTensor] = None172 prediction_logits: Optional[torch.FloatTensor] = None173 cross_relationship_score: Optional[torch.FloatTensor] = None174 question_answering_score: Optional[torch.FloatTensor] = None175 language_hidden_states: Optional[tuple[torch.FloatTensor]] = None176 vision_hidden_states: Optional[tuple[torch.FloatTensor]] = None177 language_attentions: Optional[tuple[torch.FloatTensor]] = None178 vision_attentions: Optional[tuple[torch.FloatTensor]] = None179 cross_encoder_attentions: Optional[tuple[torch.FloatTensor]] = None180 181 182def load_tf_weights_in_lxmert(model, config, tf_checkpoint_path):183 """Load tf checkpoints in a pytorch model."""184 try:185 import re186 187 import numpy as np188 import tensorflow as tf189 except ImportError:190 logger.error(191 "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "192 "https://www.tensorflow.org/install/ for installation instructions."193 )194 raise195 tf_path = os.path.abspath(tf_checkpoint_path)196 logger.info(f"Converting TensorFlow checkpoint from {tf_path}")197 # Load weights from TF model198 init_vars = tf.train.list_variables(tf_path)199 names = []200 arrays = []201 for name, shape in init_vars:202 logger.info(f"Loading TF weight {name} with shape {shape}")203 array = tf.train.load_variable(tf_path, name)204 names.append(name)205 arrays.append(array)206 207 for name, array in zip(names, arrays):208 name = name.split("/")209 # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v210 # which are not required for using pretrained model211 if any(212 n213 in [214 "adam_v",215 "adam_m",216 "AdamWeightDecayOptimizer",217 "AdamWeightDecayOptimizer_1",218 "global_step",219 ]220 for n in name221 ):222 logger.info(f"Skipping {'/'.join(name)}")223 continue224 pointer = model225 for m_name in name:226 if re.fullmatch(r"[A-Za-z]+_\d+", m_name):227 scope_names = re.split(r"_(\d+)", m_name)228 else:229 scope_names = [m_name]230 if scope_names[0] == "kernel" or scope_names[0] == "gamma":231 pointer = getattr(pointer, "weight")232 elif scope_names[0] == "output_bias" or scope_names[0] == "beta":233 pointer = getattr(pointer, "bias")234 elif scope_names[0] == "output_weights":235 pointer = getattr(pointer, "weight")236 elif scope_names[0] == "squad":237 pointer = getattr(pointer, "classifier")238 else:239 try:240 pointer = getattr(pointer, scope_names[0])241 except AttributeError:242 logger.info(f"Skipping {'/'.join(name)}")243 continue244 if len(scope_names) >= 2:245 num = int(scope_names[1])246 pointer = pointer[num]247 if m_name[-11:] == "_embeddings":248 pointer = getattr(pointer, "weight")249 elif m_name == "kernel":250 array = np.transpose(array)251 try:252 assert pointer.shape == array.shape253 except AssertionError as e:254 e.args += (pointer.shape, array.shape)255 raise256 logger.info(f"Initialize PyTorch weight {name}")257 pointer.data = torch.from_numpy(array)258 return model259 260 261class LxmertEmbeddings(nn.Module):262 """Construct the embeddings from word, position and token_type embeddings."""263 264 def __init__(self, config):265 super().__init__()266 self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0)267 self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size, padding_idx=0)268 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size, padding_idx=0)269 270 # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load271 # any TensorFlow checkpoint file272 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=1e-12)273 self.dropout = nn.Dropout(config.hidden_dropout_prob)274 275 def forward(self, input_ids, token_type_ids=None, inputs_embeds=None):276 if input_ids is not None:277 input_shape = input_ids.size()278 device = input_ids.device279 else:280 input_shape = inputs_embeds.size()[:-1]281 device = inputs_embeds.device282 seq_length = input_shape[1]283 284 position_ids = torch.arange(seq_length, dtype=torch.long, device=device)285 position_ids = position_ids.unsqueeze(0).expand(input_shape)286 287 if token_type_ids is None:288 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)289 290 if inputs_embeds is None:291 inputs_embeds = self.word_embeddings(input_ids)292 position_embeddings = self.position_embeddings(position_ids)293 token_type_embeddings = self.token_type_embeddings(token_type_ids)294 295 embeddings = inputs_embeds + position_embeddings + token_type_embeddings296 embeddings = self.LayerNorm(embeddings)297 embeddings = self.dropout(embeddings)298 return embeddings299 300 301class LxmertAttention(nn.Module):302 def __init__(self, config, ctx_dim=None):303 super().__init__()304 if config.hidden_size % config.num_attention_heads != 0:305 raise ValueError(306 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "307 f"heads ({config.num_attention_heads})"308 )309 self.num_attention_heads = config.num_attention_heads310 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)311 self.head_size = self.num_attention_heads * self.attention_head_size312 313 # visual_dim = 2048314 if ctx_dim is None:315 ctx_dim = config.hidden_size316 self.query = nn.Linear(config.hidden_size, self.head_size)317 self.key = nn.Linear(ctx_dim, self.head_size)318 self.value = nn.Linear(ctx_dim, self.head_size)319 320 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)321 322 def forward(self, hidden_states, context, attention_mask=None, output_attentions=False):323 batch_size, seq_length, _ = hidden_states.shape324 query_layer = (325 self.query(hidden_states)326 .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)327 .transpose(1, 2)328 )329 key_layer = (330 self.key(context).view(batch_size, -1, self.num_attention_heads, self.attention_head_size).transpose(1, 2)331 )332 value_layer = (333 self.value(context)334 .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)335 .transpose(1, 2)336 )337 338 # Take the dot product between "query" and "key" to get the raw attention scores.339 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))340 attention_scores = attention_scores / math.sqrt(self.attention_head_size)341 # Apply the attention mask is (precomputed for all layers in BertModel forward() function)342 if attention_mask is not None:343 attention_scores = attention_scores + attention_mask344 345 # Normalize the attention scores to probabilities.346 attention_probs = nn.functional.softmax(attention_scores, dim=-1)347 348 # This is actually dropping out entire tokens to attend to, which might349 # seem a bit unusual, but is taken from the original Transformer paper.350 attention_probs = self.dropout(attention_probs)351 352 context_layer = torch.matmul(attention_probs, value_layer)353 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()354 new_context_layer_shape = context_layer.size()[:-2] + (self.head_size,)355 context_layer = context_layer.view(new_context_layer_shape)356 357 outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)358 return outputs359 360 361class LxmertAttentionOutput(nn.Module):362 def __init__(self, config):363 super().__init__()364 self.dense = nn.Linear(config.hidden_size, config.hidden_size)365 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=1e-12)366 self.dropout = nn.Dropout(config.hidden_dropout_prob)367 368 def forward(self, hidden_states, input_tensor):369 hidden_states = self.dense(hidden_states)370 hidden_states = self.dropout(hidden_states)371 hidden_states = self.LayerNorm(hidden_states + input_tensor)372 return hidden_states373 374 375class LxmertCrossAttentionLayer(nn.Module):376 def __init__(self, config):377 super().__init__()378 self.att = LxmertAttention(config)379 self.output = LxmertAttentionOutput(config)380 381 def forward(self, input_tensor, ctx_tensor, ctx_att_mask=None, output_attentions=False):382 output = self.att(input_tensor, ctx_tensor, ctx_att_mask, output_attentions=output_attentions)383 if output_attentions:384 attention_probs = output[1]385 attention_output = self.output(output[0], input_tensor)386 outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)387 return outputs388 389 390class LxmertSelfAttentionLayer(nn.Module):391 def __init__(self, config):392 super().__init__()393 self.self = LxmertAttention(config)394 self.output = LxmertAttentionOutput(config)395 396 def forward(self, input_tensor, attention_mask, output_attentions=False):397 # Self attention attends to itself, thus keys and queries are the same (input_tensor).398 output = self.self(399 input_tensor,400 input_tensor,401 attention_mask,402 output_attentions=output_attentions,403 )404 if output_attentions:405 attention_probs = output[1]406 attention_output = self.output(output[0], input_tensor)407 outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)408 return outputs409 410 411class LxmertIntermediate(nn.Module):412 def __init__(self, config):413 super().__init__()414 self.dense = nn.Linear(config.hidden_size, config.intermediate_size)415 self.intermediate_act_fn = ACT2FN[config.hidden_act]416 417 def forward(self, hidden_states):418 hidden_states = self.dense(hidden_states)419 hidden_states = self.intermediate_act_fn(hidden_states)420 return hidden_states421 422 423class LxmertOutput(nn.Module):424 def __init__(self, config):425 super().__init__()426 self.dense = nn.Linear(config.intermediate_size, config.hidden_size)427 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=1e-12)428 self.dropout = nn.Dropout(config.hidden_dropout_prob)429 430 def forward(self, hidden_states, input_tensor):431 hidden_states = self.dense(hidden_states)432 hidden_states = self.dropout(hidden_states)433 hidden_states = self.LayerNorm(hidden_states + input_tensor)434 return hidden_states435 436 437class LxmertLayer(nn.Module):438 def __init__(self, config):439 super().__init__()440 self.attention = LxmertSelfAttentionLayer(config)441 self.intermediate = LxmertIntermediate(config)442 self.output = LxmertOutput(config)443 444 def forward(self, hidden_states, attention_mask=None, output_attentions=False):445 outputs = self.attention(hidden_states, attention_mask, output_attentions=output_attentions)446 attention_output = outputs[0]447 intermediate_output = self.intermediate(attention_output)448 layer_output = self.output(intermediate_output, attention_output)449 outputs = (layer_output,) + outputs[1:] # add attentions if we output them450 return outputs451 452 453class LxmertXLayer(nn.Module):454 def __init__(self, config):455 super().__init__()456 # The cross-attention Layer457 self.visual_attention = LxmertCrossAttentionLayer(config)458 459 # Self-attention Layers460 self.lang_self_att = LxmertSelfAttentionLayer(config)461 self.visn_self_att = LxmertSelfAttentionLayer(config)462 463 # Intermediate and Output Layers (FFNs)464 self.lang_inter = LxmertIntermediate(config)465 self.lang_output = LxmertOutput(config)466 self.visn_inter = LxmertIntermediate(config)467 self.visn_output = LxmertOutput(config)468 469 def cross_att(470 self,471 lang_input,472 lang_attention_mask,473 visual_input,474 visual_attention_mask,475 output_x_attentions=False,476 ):477 # Cross Attention478 lang_att_output = self.visual_attention(479 lang_input,480 visual_input,481 ctx_att_mask=visual_attention_mask,482 output_attentions=output_x_attentions,483 )484 visual_att_output = self.visual_attention(485 visual_input,486 lang_input,487 ctx_att_mask=lang_attention_mask,488 output_attentions=False,489 )490 return lang_att_output, visual_att_output491 492 def self_att(self, lang_input, lang_attention_mask, visual_input, visual_attention_mask):493 # Self Attention494 lang_att_output = self.lang_self_att(lang_input, lang_attention_mask, output_attentions=False)495 visual_att_output = self.visn_self_att(visual_input, visual_attention_mask, output_attentions=False)496 return lang_att_output[0], visual_att_output[0]497 498 def output_fc(self, lang_input, visual_input):499 # FC layers500 lang_inter_output = self.lang_inter(lang_input)501 visual_inter_output = self.visn_inter(visual_input)502 503 # Layer output504 lang_output = self.lang_output(lang_inter_output, lang_input)505 visual_output = self.visn_output(visual_inter_output, visual_input)506 507 return lang_output, visual_output508 509 def forward(510 self,511 lang_feats,512 lang_attention_mask,513 visual_feats,514 visual_attention_mask,515 output_attentions=False,516 ):517 lang_att_output, visual_att_output = self.cross_att(518 lang_input=lang_feats,519 lang_attention_mask=lang_attention_mask,520 visual_input=visual_feats,521 visual_attention_mask=visual_attention_mask,522 output_x_attentions=output_attentions,523 )524 attention_probs = lang_att_output[1:]525 lang_att_output, visual_att_output = self.self_att(526 lang_att_output[0],527 lang_attention_mask,528 visual_att_output[0],529 visual_attention_mask,530 )531 532 lang_output, visual_output = self.output_fc(lang_att_output, visual_att_output)533 return (534 (535 lang_output,536 visual_output,537 attention_probs[0],538 )539 if output_attentions540 else (lang_output, visual_output)541 )542 543 544class LxmertVisualFeatureEncoder(nn.Module):545 def __init__(self, config):546 super().__init__()547 feat_dim = config.visual_feat_dim548 pos_dim = config.visual_pos_dim549 550 # Object feature encoding551 self.visn_fc = nn.Linear(feat_dim, config.hidden_size)552 self.visn_layer_norm = nn.LayerNorm(config.hidden_size, eps=1e-12)553 554 # Box position encoding555 self.box_fc = nn.Linear(pos_dim, config.hidden_size)556 self.box_layer_norm = nn.LayerNorm(config.hidden_size, eps=1e-12)557 558 self.dropout = nn.Dropout(config.hidden_dropout_prob)559 560 def forward(self, visual_feats, visual_pos):561 x = self.visn_fc(visual_feats)562 x = self.visn_layer_norm(x)563 y = self.box_fc(visual_pos)564 y = self.box_layer_norm(y)565 output = (x + y) / 2566 567 output = self.dropout(output)568 return output569 570 571class LxmertEncoder(nn.Module):572 def __init__(self, config):573 super().__init__()574 575 # Obj-level image embedding layer576 self.visn_fc = LxmertVisualFeatureEncoder(config)577 self.config = config578 579 # Number of layers580 self.num_l_layers = config.l_layers581 self.num_x_layers = config.x_layers582 self.num_r_layers = config.r_layers583 584 # Layers585 # Using self.layer instead of self.l_layer to support loading BERT weights.586 self.layer = nn.ModuleList([LxmertLayer(config) for _ in range(self.num_l_layers)])587 self.x_layers = nn.ModuleList([LxmertXLayer(config) for _ in range(self.num_x_layers)])588 self.r_layers = nn.ModuleList([LxmertLayer(config) for _ in range(self.num_r_layers)])589 590 def forward(591 self,592 lang_feats,593 lang_attention_mask,594 visual_feats,595 visual_pos,596 visual_attention_mask=None,597 output_attentions=None,598 ):599 vision_hidden_states = ()600 language_hidden_states = ()601 vision_attentions = () if output_attentions or self.config.output_attentions else None602 language_attentions = () if output_attentions or self.config.output_attentions else None603 cross_encoder_attentions = () if output_attentions or self.config.output_attentions else None604 605 visual_feats = self.visn_fc(visual_feats, visual_pos)606 607 # Run language layers608 for layer_module in self.layer:609 l_outputs = layer_module(lang_feats, lang_attention_mask, output_attentions=output_attentions)610 lang_feats = l_outputs[0]611 language_hidden_states = language_hidden_states + (lang_feats,)612 if language_attentions is not None:613 language_attentions = language_attentions + (l_outputs[1],)614 615 # Run relational layers616 for layer_module in self.r_layers:617 v_outputs = layer_module(visual_feats, visual_attention_mask, output_attentions=output_attentions)618 visual_feats = v_outputs[0]619 vision_hidden_states = vision_hidden_states + (visual_feats,)620 if vision_attentions is not None:621 vision_attentions = vision_attentions + (v_outputs[1],)622 623 # Run cross-modality layers624 for layer_module in self.x_layers:625 x_outputs = layer_module(626 lang_feats,627 lang_attention_mask,628 visual_feats,629 visual_attention_mask,630 output_attentions=output_attentions,631 )632 lang_feats, visual_feats = x_outputs[:2]633 vision_hidden_states = vision_hidden_states + (visual_feats,)634 language_hidden_states = language_hidden_states + (lang_feats,)635 if cross_encoder_attentions is not None:636 cross_encoder_attentions = cross_encoder_attentions + (x_outputs[2],)637 visual_encoder_outputs = (638 vision_hidden_states,639 vision_attentions if output_attentions else None,640 )641 lang_encoder_outputs = (642 language_hidden_states,643 language_attentions if output_attentions else None,644 )645 return (646 visual_encoder_outputs,647 lang_encoder_outputs,648 cross_encoder_attentions if output_attentions else None,649 )650 651 652class LxmertPooler(nn.Module):653 def __init__(self, config):654 super().__init__()655 self.dense = nn.Linear(config.hidden_size, config.hidden_size)656 self.activation = nn.Tanh()657 658 def forward(self, hidden_states):659 # We "pool" the model by simply taking the hidden state corresponding660 # to the first token.661 first_token_tensor = hidden_states[:, 0]662 pooled_output = self.dense(first_token_tensor)663 pooled_output = self.activation(pooled_output)664 return pooled_output665 666 667class LxmertPredictionHeadTransform(nn.Module):668 def __init__(self, config):669 super().__init__()670 self.dense = nn.Linear(config.hidden_size, config.hidden_size)671 self.transform_act_fn = ACT2FN[config.hidden_act]672 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=1e-12)673 674 def forward(self, hidden_states):675 hidden_states = self.dense(hidden_states)676 hidden_states = self.transform_act_fn(hidden_states)677 hidden_states = self.LayerNorm(hidden_states)678 return hidden_states679 680 681class LxmertLMPredictionHead(nn.Module):682 def __init__(self, config, lxmert_model_embedding_weights):683 super().__init__()684 self.transform = LxmertPredictionHeadTransform(config)685 686 # The output weights are the same as the input embeddings, but there is687 # an output-only bias for each token.688 self.decoder = nn.Linear(689 lxmert_model_embedding_weights.size(1),690 lxmert_model_embedding_weights.size(0),691 bias=False,692 )693 self.decoder.weight = lxmert_model_embedding_weights694 self.bias = nn.Parameter(torch.zeros(lxmert_model_embedding_weights.size(0)))695 696 def forward(self, hidden_states):697 hidden_states = self.transform(hidden_states)698 hidden_states = self.decoder(hidden_states) + self.bias699 return hidden_states700 701 702class LxmertVisualAnswerHead(nn.Module):703 def __init__(self, config, num_labels):704 super().__init__()705 hid_dim = config.hidden_size706 self.logit_fc = nn.Sequential(707 nn.Linear(hid_dim, hid_dim * 2),708 GeLU(),709 nn.LayerNorm(hid_dim * 2, eps=1e-12),710 nn.Linear(hid_dim * 2, num_labels),711 )712 713 def forward(self, hidden_states):714 return self.logit_fc(hidden_states)715 716 717class LxmertVisualObjHead(nn.Module):718 def __init__(self, config):719 super().__init__()720 self.transform = LxmertPredictionHeadTransform(config)721 # Decide the use of visual losses722 visual_losses = {}723 if config.visual_obj_loss:724 visual_losses["obj"] = {"shape": (-1,), "num": config.num_object_labels}725 if config.visual_attr_loss:726 visual_losses["attr"] = {"shape": (-1,), "num": config.num_attr_labels}727 if config.visual_feat_loss:728 visual_losses["feat"] = {729 "shape": (-1, config.visual_feat_dim),730 "num": config.visual_feat_dim,731 }732 self.visual_losses = visual_losses733 734 # The output weights are the same as the input embeddings, but there is735 # an output-only bias for each token.736 self.decoder_dict = nn.ModuleDict(737 {key: nn.Linear(config.hidden_size, self.visual_losses[key]["num"]) for key in self.visual_losses}738 )739 740 def forward(self, hidden_states):741 hidden_states = self.transform(hidden_states)742 output = {}743 for key in self.visual_losses:744 output[key] = self.decoder_dict[key](hidden_states)745 return output746 747 748class LxmertPreTrainingHeads(nn.Module):749 def __init__(self, config, lxmert_model_embedding_weights):750 super().__init__()751 self.predictions = LxmertLMPredictionHead(config, lxmert_model_embedding_weights)752 self.seq_relationship = nn.Linear(config.hidden_size, 2)753 754 def forward(self, sequence_output, pooled_output):755 prediction_scores = self.predictions(sequence_output)756 seq_relationship_score = self.seq_relationship(pooled_output)757 return prediction_scores, seq_relationship_score758 759 760@auto_docstring761class LxmertPreTrainedModel(PreTrainedModel):762 config: LxmertConfig763 load_tf_weights = load_tf_weights_in_lxmert764 base_model_prefix = "lxmert"765 _supports_param_buffer_assignment = False766 767 def _init_weights(self, module):768 """Initialize the weights"""769 if isinstance(module, nn.Linear):770 # Slightly different from the TF version which uses truncated_normal for initialization771 # cf https://github.com/pytorch/pytorch/pull/5617772 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)773 if module.bias is not None:774 module.bias.data.zero_()775 elif isinstance(module, nn.Embedding):776 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)777 if module.padding_idx is not None:778 module.weight.data[module.padding_idx].zero_()779 elif isinstance(module, nn.LayerNorm):780 module.bias.data.zero_()781 module.weight.data.fill_(1.0)782 elif isinstance(module, LxmertLMPredictionHead):783 module.bias.data.zero_()784 785 786@auto_docstring787class LxmertModel(LxmertPreTrainedModel):788 def __init__(self, config):789 super().__init__(config)790 self.embeddings = LxmertEmbeddings(config)791 self.encoder = LxmertEncoder(config)792 self.pooler = LxmertPooler(config)793 # Initialize weights and apply final processing794 self.post_init()795 796 def get_input_embeddings(self):797 return self.embeddings.word_embeddings798 799 def set_input_embeddings(self, new_embeddings):800 self.embeddings.word_embeddings = new_embeddings801 802 @auto_docstring803 def forward(804 self,805 input_ids: Optional[torch.LongTensor] = None,806 visual_feats: Optional[torch.FloatTensor] = None,807 visual_pos: Optional[torch.FloatTensor] = None,808 attention_mask: Optional[torch.FloatTensor] = None,809 visual_attention_mask: Optional[torch.FloatTensor] = None,810 token_type_ids: Optional[torch.LongTensor] = None,811 inputs_embeds: Optional[torch.FloatTensor] = None,812 output_attentions: Optional[bool] = None,813 output_hidden_states: Optional[bool] = None,814 return_dict: Optional[bool] = None,815 ) -> Union[LxmertModelOutput, tuple[torch.FloatTensor]]:816 r"""817 visual_feats (`torch.FloatTensor` of shape `(batch_size, num_visual_features, visual_feat_dim)`):818 This input represents visual features. They ROI pooled object features from bounding boxes using a819 faster-RCNN model)820 821 These are currently not provided by the transformers library.822 visual_pos (`torch.FloatTensor` of shape `(batch_size, num_visual_features, visual_pos_dim)`):823 This input represents spatial features corresponding to their relative (via index) visual features. The824 pre-trained LXMERT model expects these spatial features to be normalized bounding boxes on a scale of 0 to825 1.826 827 These are currently not provided by the transformers library.828 visual_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):829 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:830 831 - 1 for tokens that are **not masked**,832 - 0 for tokens that are **masked**.833 834 [What are attention masks?](../glossary#attention-mask)835 """836 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions837 output_hidden_states = (838 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states839 )840 return_dict = return_dict if return_dict is not None else self.config.use_return_dict841 842 if input_ids is not None and inputs_embeds is not None:843 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")844 elif input_ids is not None:845 self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)846 input_shape = input_ids.size()847 elif inputs_embeds is not None:848 input_shape = inputs_embeds.size()[:-1]849 else:850 raise ValueError("You have to specify either input_ids or inputs_embeds")851 852 if visual_feats is None:853 raise ValueError("`visual_feats` cannot be `None`")854 if visual_pos is None:855 raise ValueError("`visual_pos` cannot be `None`")856 857 device = input_ids.device if input_ids is not None else inputs_embeds.device858 859 if attention_mask is None:860 attention_mask = torch.ones(input_shape, device=device)861 if token_type_ids is None:862 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)863 864 # We create a 3D attention mask from a 2D tensor mask.865 # Sizes are [batch_size, 1, 1, to_seq_length]866 # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]867 # this attention mask is more simple than the triangular masking of causal attention868 # used in OpenAI GPT, we just need to prepare the broadcast dimension here.869 extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)870 871 # Since attention_mask is 1.0 for positions we want to attend and 0.0 for872 # masked positions, this operation will create a tensor which is 0.0 for873 # positions we want to attend and the dtype's smallest value for masked positions.874 # Since we are adding it to the raw scores before the softmax, this is875 # effectively the same as removing these entirely.876 extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)877 extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(self.dtype).min878 879 # Process the visual attention mask880 if visual_attention_mask is not None:881 extended_visual_attention_mask = visual_attention_mask.unsqueeze(1).unsqueeze(2)882 extended_visual_attention_mask = extended_visual_attention_mask.to(dtype=self.dtype)883 extended_visual_attention_mask = (1.0 - extended_visual_attention_mask) * torch.finfo(self.dtype).min884 else:885 extended_visual_attention_mask = None886 887 # Positional Word Embeddings888 embedding_output = self.embeddings(input_ids, token_type_ids, inputs_embeds)889 890 # Run Lxmert encoder891 encoder_outputs = self.encoder(892 embedding_output,893 extended_attention_mask,894 visual_feats=visual_feats,895 visual_pos=visual_pos,896 visual_attention_mask=extended_visual_attention_mask,897 output_attentions=output_attentions,898 )899 900 visual_encoder_outputs, lang_encoder_outputs = encoder_outputs[:2]901 vision_hidden_states = visual_encoder_outputs[0]902 language_hidden_states = lang_encoder_outputs[0]903 904 all_attentions = ()905 if output_attentions:906 language_attentions = lang_encoder_outputs[1]907 vision_attentions = visual_encoder_outputs[1]908 cross_encoder_attentions = encoder_outputs[2]909 all_attentions = (910 language_attentions,911 vision_attentions,912 cross_encoder_attentions,913 )914 915 hidden_states = (language_hidden_states, vision_hidden_states) if output_hidden_states else ()916 917 visual_output = vision_hidden_states[-1]918 lang_output = language_hidden_states[-1]919 pooled_output = self.pooler(lang_output)920 921 if not return_dict:922 return (lang_output, visual_output, pooled_output) + hidden_states + all_attentions923 924 return LxmertModelOutput(925 pooled_output=pooled_output,926 language_output=lang_output,927 vision_output=visual_output,928 language_hidden_states=language_hidden_states if output_hidden_states else None,929 vision_hidden_states=vision_hidden_states if output_hidden_states else None,930 language_attentions=language_attentions if output_attentions else None,931 vision_attentions=vision_attentions if output_attentions else None,932 cross_encoder_attentions=cross_encoder_attentions if output_attentions else None,933 )934 935 936@auto_docstring937class LxmertForPreTraining(LxmertPreTrainedModel):938 _tied_weights_keys = ["cls.predictions.decoder.weight"]939 940 def __init__(self, config):941 super().__init__(config)942 # Configuration943 self.config = config944 self.num_qa_labels = config.num_qa_labels945 self.visual_loss_normalizer = config.visual_loss_normalizer946 947 # Use of pretraining tasks948 self.task_mask_lm = config.task_mask_lm949 self.task_obj_predict = config.task_obj_predict950 self.task_matched = config.task_matched951 self.task_qa = config.task_qa952 953 # Lxmert backbone954 self.lxmert = LxmertModel(config)955 956 # Pre-training heads957 self.cls = LxmertPreTrainingHeads(config, self.lxmert.embeddings.word_embeddings.weight)958 if self.task_obj_predict:959 self.obj_predict_head = LxmertVisualObjHead(config)960 if self.task_qa:961 self.answer_head = LxmertVisualAnswerHead(config, self.num_qa_labels)962 963 # Weight initialization964 # Initialize weights and apply final processing965 self.post_init()966 967 # Loss functions968 self.loss_fcts = {969 "l2": SmoothL1Loss(reduction="none"),970 "visual_ce": CrossEntropyLoss(reduction="none"),971 "ce": CrossEntropyLoss(),972 }973 974 visual_losses = {}975 if config.visual_obj_loss:976 visual_losses["obj"] = {977 "shape": (-1,),978 "num": config.num_object_labels,979 "loss": "visual_ce",980 }981 if config.visual_attr_loss:982 visual_losses["attr"] = {983 "shape": (-1,),984 "num": config.num_attr_labels,985 "loss": "visual_ce",986 }987 if config.visual_feat_loss:988 visual_losses["feat"] = {989 "shape": (-1, config.visual_feat_dim),990 "num": config.visual_feat_dim,991 "loss": "l2",992 }993 self.visual_losses = visual_losses994 995 def _tie_weights(self):996 self.cls.predictions.decoder.weight = self.lxmert.embeddings.word_embeddings.weight997 998 def resize_token_embeddings(999 self, new_num_tokens: int, pad_to_multiple_of: Optional[int] = None, mean_resizing: bool = True1000 ) -> nn.Embedding:1001 # Adding the following steps to resize bias to match the shape of resized embeddings1002 new_embeddings = super().resize_token_embeddings(new_num_tokens, pad_to_multiple_of, mean_resizing)1003 self.cls.predictions.bias = self._resize_bias(self.cls.predictions.bias, new_num_tokens)1004 return new_embeddings1005 1006 def _resize_bias(self, bias, new_num_tokens: int):1007 old_num_tokens = bias.shape[0]1008 if new_num_tokens <= old_num_tokens:1009 new_bias = bias[:new_num_tokens]1010 else:1011 extra_bias = torch.zeros(new_num_tokens - old_num_tokens, device=bias.device)1012 new_bias = torch.cat([bias, extra_bias])1013 new_bias = nn.Parameter(new_bias)1014 return new_bias1015 1016 def resize_num_qa_labels(self, num_labels):1017 """1018 Build a resized question answering linear layer Module from a provided new linear layer. Increasing the size1019 will add newly initialized weights. Reducing the size will remove weights from the end1020 1021 Args:1022 num_labels (`int`, *optional*):1023 New number of labels in the linear layer weight matrix. Increasing the size will add newly initialized1024 weights at the end. Reducing the size will remove weights from the end. If not provided or `None`, just1025 returns a pointer to the qa labels ``torch.nn.Linear``` module of the model without doing anything.1026 1027 Return:1028 `torch.nn.Linear`: Pointer to the resized Linear layer or the old Linear layer1029 """1030 1031 cur_qa_logit_layer = self.get_qa_logit_layer()1032 if num_labels is None or cur_qa_logit_layer is None:1033 return1034 new_qa_logit_layer = self._resize_qa_labels(num_labels)1035 self.config.num_qa_labels = num_labels1036 self.num_qa_labels = num_labels1037 1038 return new_qa_logit_layer1039 1040 def _resize_qa_labels(self, num_labels):1041 cur_qa_logit_layer = self.get_qa_logit_layer()1042 new_qa_logit_layer = self._get_resized_qa_labels(cur_qa_logit_layer, num_labels)1043 self._set_qa_logit_layer(new_qa_logit_layer)1044 return self.get_qa_logit_layer()1045 1046 def get_qa_logit_layer(self) -> nn.Module:1047 """1048 Returns the linear layer that produces question answering logits.1049 1050 Returns:1051 `nn.Module`: A torch module mapping the question answering prediction hidden states or `None` if LXMERT1052 does not have a visual answering head.1053 """1054 if hasattr(self, "answer_head"):1055 return self.answer_head.logit_fc[-1]1056 1057 def _set_qa_logit_layer(self, qa_logit_layer):1058 self.answer_head.logit_fc[-1] = qa_logit_layer1059 1060 def _get_resized_qa_labels(self, cur_qa_logit_layer, num_labels):1061 if num_labels is None:1062 return cur_qa_logit_layer1063 1064 cur_qa_labels, hidden_dim = cur_qa_logit_layer.weight.size()1065 if cur_qa_labels == num_labels:1066 return cur_qa_logit_layer1067 1068 # Build new linear output1069 if getattr(cur_qa_logit_layer, "bias", None) is not None:1070 new_qa_logit_layer = nn.Linear(hidden_dim, num_labels)1071 else:1072 new_qa_logit_layer = nn.Linear(hidden_dim, num_labels, bias=False)1073 1074 new_qa_logit_layer.to(cur_qa_logit_layer.weight.device)1075 1076 # initialize all new labels1077 self._init_weights(new_qa_logit_layer)1078 1079 # Copy labels from the previous weights1080 num_labels_to_copy = min(cur_qa_labels, num_labels)1081 new_qa_logit_layer.weight.data[:num_labels_to_copy, :] = cur_qa_logit_layer.weight.data[:num_labels_to_copy, :]1082 if getattr(cur_qa_logit_layer, "bias", None) is not None:1083 new_qa_logit_layer.bias.data[:num_labels_to_copy] = cur_qa_logit_layer.bias.data[:num_labels_to_copy]1084 1085 return new_qa_logit_layer1086 1087 @auto_docstring1088 def forward(1089 self,1090 input_ids: Optional[torch.LongTensor] = None,1091 visual_feats: Optional[torch.FloatTensor] = None,1092 visual_pos: Optional[torch.FloatTensor] = None,1093 attention_mask: Optional[torch.FloatTensor] = None,1094 visual_attention_mask: Optional[torch.FloatTensor] = None,1095 token_type_ids: Optional[torch.LongTensor] = None,1096 inputs_embeds: Optional[torch.FloatTensor] = None,1097 labels: Optional[torch.LongTensor] = None,1098 obj_labels: Optional[dict[str, tuple[torch.FloatTensor, torch.FloatTensor]]] = None,1099 matched_label: Optional[torch.LongTensor] = None,1100 ans: Optional[torch.Tensor] = None,1101 output_attentions: Optional[bool] = None,1102 output_hidden_states: Optional[bool] = None,1103 return_dict: Optional[bool] = None,1104 **kwargs,1105 ) -> Union[LxmertForPreTrainingOutput, tuple[torch.FloatTensor]]:1106 r"""1107 visual_feats (`torch.FloatTensor` of shape `(batch_size, num_visual_features, visual_feat_dim)`):1108 This input represents visual features. They ROI pooled object features from bounding boxes using a1109 faster-RCNN model)1110 1111 These are currently not provided by the transformers library.1112 visual_pos (`torch.FloatTensor` of shape `(batch_size, num_visual_features, visual_pos_dim)`):1113 This input represents spatial features corresponding to their relative (via index) visual features. The1114 pre-trained LXMERT model expects these spatial features to be normalized bounding boxes on a scale of 0 to1115 1.1116 1117 These are currently not provided by the transformers library.1118 visual_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):1119 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:1120 1121 - 1 for tokens that are **not masked**,1122 - 0 for tokens that are **masked**.1123 1124 [What are attention masks?](../glossary#attention-mask)1125 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1126 Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,1127 config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the1128 loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`1129 obj_labels (`dict[Str: tuple[Torch.FloatTensor, Torch.FloatTensor]]`, *optional*):1130 each key is named after each one of the visual losses and each element of the tuple is of the shape1131 `(batch_size, num_features)` and `(batch_size, num_features, visual_feature_dim)` for each the label id and1132 the label score respectively1133 matched_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1134 Labels for computing the whether or not the text input matches the image (classification) loss. Input1135 should be a sequence pair (see `input_ids` docstring) Indices should be in `[0, 1]`:1136 1137 - 0 indicates that the sentence does not match the image,1138 - 1 indicates that the sentence does match the image.1139 ans (`Torch.Tensor` of shape `(batch_size)`, *optional*):1140 a one hot representation hof the correct answer *optional*1141 """1142 1143 if "masked_lm_labels" in kwargs:1144 warnings.warn(1145 "The `masked_lm_labels` argument is deprecated and will be removed in a future version, use `labels`"1146 " instead.",1147 FutureWarning,1148 )1149 labels = kwargs.pop("masked_lm_labels")1150 1151 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1152 1153 device = input_ids.device if input_ids is not None else inputs_embeds.device1154 lxmert_output = self.lxmert(1155 input_ids=input_ids,1156 visual_feats=visual_feats,1157 visual_pos=visual_pos,1158 token_type_ids=token_type_ids,1159 attention_mask=attention_mask,1160 visual_attention_mask=visual_attention_mask,1161 inputs_embeds=inputs_embeds,1162 output_hidden_states=output_hidden_states,1163 output_attentions=output_attentions,1164 return_dict=return_dict,1165 )1166 1167 lang_output, visual_output, pooled_output = (1168 lxmert_output[0],1169 lxmert_output[1],1170 lxmert_output[2],1171 )1172 lang_prediction_scores, cross_relationship_score = self.cls(lang_output, pooled_output)1173 if self.task_qa:1174 answer_score = self.answer_head(pooled_output)1175 else:1176 answer_score = pooled_output[0][0]1177 1178 total_loss = (1179 None1180 if (labels is None and matched_label is None and obj_labels is None and ans is None)1181 else torch.tensor(0.0, device=device)1182 )1183 if labels is not None and self.task_mask_lm:1184 masked_lm_loss = self.loss_fcts["ce"](1185 lang_prediction_scores.view(-1, self.config.vocab_size),1186 labels.view(-1),1187 )1188 total_loss += masked_lm_loss1189 if matched_label is not None and self.task_matched:1190 matched_loss = self.loss_fcts["ce"](cross_relationship_score.view(-1, 2), matched_label.view(-1))1191 total_loss += matched_loss1192 if obj_labels is not None and self.task_obj_predict:1193 total_visual_loss = torch.tensor(0.0, device=input_ids.device)1194 visual_prediction_scores_dict = self.obj_predict_head(visual_output)1195 for key, key_info in self.visual_losses.items():1196 label, mask_conf = obj_labels[key]1197 output_dim = key_info["num"]1198 loss_fct_name = key_info["loss"]1199 label_shape = key_info["shape"]1200 weight = self.visual_loss_normalizer