DoruC/Grounded-Segment-Anything
0
1# coding=utf-82# Copyright 2022 Meta Platforms 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 Data2VecVision model."""16 17 18import collections.abc19import math20from dataclasses import dataclass21from typing import List, Optional, Tuple, Union22 23import torch24import torch.utils.checkpoint25from torch import nn26from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss27 28from ...activations import ACT2FN29from ...modeling_outputs import (30 BaseModelOutput,31 BaseModelOutputWithPooling,32 ImageClassifierOutput,33 SemanticSegmenterOutput,34)35from ...modeling_utils import PreTrainedModel36from ...pytorch_utils import find_pruneable_heads_and_indices, meshgrid, prune_linear_layer37from ...utils import (38 add_code_sample_docstrings,39 add_start_docstrings,40 add_start_docstrings_to_model_forward,41 logging,42 replace_return_docstrings,43)44from .configuration_data2vec_vision import Data2VecVisionConfig45 46 47logger = logging.get_logger(__name__)48 49# General docstring50_CONFIG_FOR_DOC = "Data2VecVisionConfig"51 52# Base docstring53_CHECKPOINT_FOR_DOC = "facebook/data2vec-vision-base"54_EXPECTED_OUTPUT_SHAPE = [1, 197, 768]55 56# Image classification docstring57_IMAGE_CLASS_CHECKPOINT = "facebook/data2vec-vision-base-ft1k"58_IMAGE_CLASS_EXPECTED_OUTPUT = "remote control, remote"59 60DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST = [61 "facebook/data2vec-vision-base-ft1k",62 # See all Data2VecVision models at https://huggingface.co/models?filter=data2vec-vision63]64 65 66@dataclass67# Copied from transformers.models.beit.modeling_beit.BeitModelOutputWithPooling with Beit->Data2VecVision68class Data2VecVisionModelOutputWithPooling(BaseModelOutputWithPooling):69 """70 Class for outputs of [`Data2VecVisionModel`].71 72 Args:73 last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):74 Sequence of hidden-states at the output of the last layer of the model.75 pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`):76 Average of the last layer hidden states of the patch tokens (excluding the *[CLS]* token) if77 *config.use_mean_pooling* is set to True. If set to False, then the final hidden state of the *[CLS]* token78 will be returned.79 hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):80 Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of81 shape `(batch_size, sequence_length, hidden_size)`.82 83 Hidden-states of the model at the output of each layer plus the initial embedding outputs.84 attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):85 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,86 sequence_length)`.87 88 Attentions weights after the attention softmax, used to compute the weighted average in the self-attention89 heads.90 """91 92 93# Copied from transformers.models.beit.modeling_beit.drop_path94def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:95 """96 Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).97 98 Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,99 however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...100 See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the101 layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the102 argument.103 """104 if drop_prob == 0.0 or not training:105 return input106 keep_prob = 1 - drop_prob107 shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets108 random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)109 random_tensor.floor_() # binarize110 output = input.div(keep_prob) * random_tensor111 return output112 113 114# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->Data2VecVision115class Data2VecVisionDropPath(nn.Module):116 """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""117 118 def __init__(self, drop_prob: Optional[float] = None) -> None:119 super().__init__()120 self.drop_prob = drop_prob121 122 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:123 return drop_path(hidden_states, self.drop_prob, self.training)124 125 def extra_repr(self) -> str:126 return "p={}".format(self.drop_prob)127 128 129# Copied from transformers.models.beit.modeling_beit.BeitEmbeddings with Beit->Data2VecVision130class Data2VecVisionEmbeddings(nn.Module):131 """132 Construct the CLS token, position and patch embeddings. Optionally, also the mask token.133 134 """135 136 def __init__(self, config: Data2VecVisionConfig) -> None:137 super().__init__()138 139 self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))140 if config.use_mask_token:141 self.mask_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))142 else:143 self.mask_token = None144 self.patch_embeddings = Data2VecVisionPatchEmbeddings(config)145 num_patches = self.patch_embeddings.num_patches146 if config.use_absolute_position_embeddings:147 self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size))148 else:149 self.position_embeddings = None150 self.dropout = nn.Dropout(config.hidden_dropout_prob)151 152 def forward(self, pixel_values: torch.Tensor, bool_masked_pos: Optional[torch.BoolTensor] = None) -> torch.Tensor:153 embeddings = self.patch_embeddings(pixel_values)154 batch_size, seq_len, _ = embeddings.size()155 156 cls_tokens = self.cls_token.expand(batch_size, -1, -1)157 if bool_masked_pos is not None:158 mask_tokens = self.mask_token.expand(batch_size, seq_len, -1)159 # replace the masked visual tokens by mask_tokens160 w = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens)161 embeddings = embeddings * (1 - w) + mask_tokens * w162 163 embeddings = torch.cat((cls_tokens, embeddings), dim=1)164 if self.position_embeddings is not None:165 embeddings = embeddings + self.position_embeddings166 embeddings = self.dropout(embeddings)167 168 return embeddings169 170 171# Copied from transformers.models.beit.modeling_beit.BeitPatchEmbeddings with Beit->Data2VecVision172class Data2VecVisionPatchEmbeddings(nn.Module):173 """174 This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial175 `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a176 Transformer.177 """178 179 def __init__(self, config):180 super().__init__()181 image_size, patch_size = config.image_size, config.patch_size182 num_channels, hidden_size = config.num_channels, config.hidden_size183 184 image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)185 patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)186 num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])187 patch_shape = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])188 self.image_size = image_size189 self.patch_size = patch_size190 self.num_channels = num_channels191 self.num_patches = num_patches192 self.patch_shape = patch_shape193 194 self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)195 196 def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:197 batch_size, num_channels, height, width = pixel_values.shape198 if num_channels != self.num_channels:199 raise ValueError(200 "Make sure that the channel dimension of the pixel values match with the one set in the configuration."201 )202 if height != self.image_size[0] or width != self.image_size[1]:203 raise ValueError(204 f"Input image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})."205 )206 embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2)207 208 return embeddings209 210 211# Copied from transformers.models.beit.modeling_beit.BeitSelfAttention with Beit->Data2VecVision212class Data2VecVisionSelfAttention(nn.Module):213 def __init__(self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None) -> None:214 super().__init__()215 if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):216 raise ValueError(217 f"The hidden size {config.hidden_size,} is not a multiple of the number of attention "218 f"heads {config.num_attention_heads}."219 )220 221 self.num_attention_heads = config.num_attention_heads222 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)223 self.all_head_size = self.num_attention_heads * self.attention_head_size224 225 self.query = nn.Linear(config.hidden_size, self.all_head_size)226 self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=False)227 self.value = nn.Linear(config.hidden_size, self.all_head_size)228 229 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)230 231 if window_size:232 self.relative_position_bias = Data2VecVisionRelativePositionBias(config, window_size=window_size)233 else:234 self.relative_position_bias = None235 236 def transpose_for_scores(self, x):237 new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)238 x = x.view(*new_x_shape)239 return x.permute(0, 2, 1, 3)240 241 def forward(242 self,243 hidden_states: torch.Tensor,244 head_mask: Optional[torch.Tensor] = None,245 output_attentions: bool = False,246 relative_position_bias: Optional["Data2VecVisionRelativePositionBias"] = None,247 ) -> Union[Tuple[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:248 mixed_query_layer = self.query(hidden_states)249 250 key_layer = self.transpose_for_scores(self.key(hidden_states))251 value_layer = self.transpose_for_scores(self.value(hidden_states))252 query_layer = self.transpose_for_scores(mixed_query_layer)253 254 # Take the dot product between "query" and "key" to get the raw attention scores.255 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))256 257 attention_scores = attention_scores / math.sqrt(self.attention_head_size)258 259 # Add relative position bias if present.260 if self.relative_position_bias is not None:261 attention_scores = attention_scores + self.relative_position_bias().unsqueeze(0)262 263 # Add shared relative position bias if provided.264 if relative_position_bias is not None:265 attention_scores = attention_scores + relative_position_bias266 267 # Normalize the attention scores to probabilities.268 attention_probs = nn.functional.softmax(attention_scores, dim=-1)269 270 # This is actually dropping out entire tokens to attend to, which might271 # seem a bit unusual, but is taken from the original Transformer paper.272 attention_probs = self.dropout(attention_probs)273 274 # Mask heads if we want to275 if head_mask is not None:276 attention_probs = attention_probs * head_mask277 278 context_layer = torch.matmul(attention_probs, value_layer)279 280 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()281 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)282 context_layer = context_layer.view(*new_context_layer_shape)283 284 outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)285 286 return outputs287 288 289# Copied from transformers.models.beit.modeling_beit.BeitSelfOutput with Beit->Data2VecVision290class Data2VecVisionSelfOutput(nn.Module):291 """292 The residual connection is defined in Data2VecVisionLayer instead of here (as is the case with other models), due293 to the layernorm applied before each block.294 """295 296 def __init__(self, config: Data2VecVisionConfig) -> None:297 super().__init__()298 self.dense = nn.Linear(config.hidden_size, config.hidden_size)299 self.dropout = nn.Dropout(config.hidden_dropout_prob)300 301 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor, gamma=None) -> torch.Tensor:302 hidden_states = self.dense(hidden_states)303 hidden_states = self.dropout(hidden_states)304 305 return hidden_states306 307 308# Copied from transformers.models.beit.modeling_beit.BeitAttention with Beit->Data2VecVision309class Data2VecVisionAttention(nn.Module):310 def __init__(self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None) -> None:311 super().__init__()312 self.attention = Data2VecVisionSelfAttention(config, window_size=window_size)313 self.output = Data2VecVisionSelfOutput(config)314 self.pruned_heads = set()315 316 def prune_heads(self, heads):317 if len(heads) == 0:318 return319 heads, index = find_pruneable_heads_and_indices(320 heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads321 )322 323 # Prune linear layers324 self.attention.query = prune_linear_layer(self.attention.query, index)325 self.attention.key = prune_linear_layer(self.attention.key, index)326 self.attention.value = prune_linear_layer(self.attention.value, index)327 self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)328 329 # Update hyper params and store pruned heads330 self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)331 self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads332 self.pruned_heads = self.pruned_heads.union(heads)333 334 def forward(335 self,336 hidden_states: torch.Tensor,337 head_mask: Optional[torch.Tensor] = None,338 output_attentions: bool = False,339 relative_position_bias: Optional["Data2VecVisionRelativePositionBias"] = None,340 ) -> Union[Tuple[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:341 self_outputs = self.attention(hidden_states, head_mask, output_attentions, relative_position_bias)342 343 attention_output = self.output(self_outputs[0], hidden_states)344 345 outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them346 return outputs347 348 349# Copied from transformers.models.beit.modeling_beit.BeitIntermediate with Beit->Data2VecVision350class Data2VecVisionIntermediate(nn.Module):351 def __init__(self, config: Data2VecVisionConfig) -> None:352 super().__init__()353 self.dense = nn.Linear(config.hidden_size, config.intermediate_size)354 if isinstance(config.hidden_act, str):355 self.intermediate_act_fn = ACT2FN[config.hidden_act]356 else:357 self.intermediate_act_fn = config.hidden_act358 359 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:360 hidden_states = self.dense(hidden_states)361 hidden_states = self.intermediate_act_fn(hidden_states)362 363 return hidden_states364 365 366# Copied from transformers.models.beit.modeling_beit.BeitOutput with Beit->Data2VecVision367class Data2VecVisionOutput(nn.Module):368 def __init__(self, config: Data2VecVisionConfig) -> None:369 super().__init__()370 self.dense = nn.Linear(config.intermediate_size, config.hidden_size)371 self.dropout = nn.Dropout(config.hidden_dropout_prob)372 373 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:374 hidden_states = self.dense(hidden_states)375 hidden_states = self.dropout(hidden_states)376 377 return hidden_states378 379 380# Copied from transformers.models.beit.modeling_beit.BeitLayer with Beit->Data2VecVision,BEiT->Data2VecVision381class Data2VecVisionLayer(nn.Module):382 """This corresponds to the Block class in the timm implementation."""383 384 def __init__(385 self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None, drop_path_rate: float = 0.0386 ) -> None:387 super().__init__()388 self.chunk_size_feed_forward = config.chunk_size_feed_forward389 self.seq_len_dim = 1390 self.attention = Data2VecVisionAttention(config, window_size=window_size)391 self.intermediate = Data2VecVisionIntermediate(config)392 self.output = Data2VecVisionOutput(config)393 self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)394 self.drop_path = Data2VecVisionDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()395 self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)396 397 init_values = config.layer_scale_init_value398 if init_values > 0:399 self.lambda_1 = nn.Parameter(init_values * torch.ones((config.hidden_size)), requires_grad=True)400 self.lambda_2 = nn.Parameter(init_values * torch.ones((config.hidden_size)), requires_grad=True)401 else:402 self.lambda_1, self.lambda_2 = None, None403 404 def forward(405 self,406 hidden_states: torch.Tensor,407 head_mask: Optional[torch.Tensor] = None,408 output_attentions: bool = False,409 relative_position_bias: Optional["Data2VecVisionRelativePositionBias"] = None,410 ) -> Union[Tuple[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:411 self_attention_outputs = self.attention(412 self.layernorm_before(hidden_states), # in Data2VecVision, layernorm is applied before self-attention413 head_mask,414 output_attentions=output_attentions,415 relative_position_bias=relative_position_bias,416 )417 attention_output = self_attention_outputs[0]418 outputs = self_attention_outputs[1:] # add self attentions if we output attention weights419 420 # apply lambda_1 if present421 if self.lambda_1 is not None:422 attention_output = self.lambda_1 * attention_output423 424 # first residual connection425 hidden_states = self.drop_path(attention_output) + hidden_states426 427 # in Data2VecVision, layernorm is also applied after self-attention428 layer_output = self.layernorm_after(hidden_states)429 430 layer_output = self.intermediate(layer_output)431 layer_output = self.output(layer_output)432 433 if self.lambda_2 is not None:434 layer_output = self.lambda_2 * layer_output435 436 # second residual connection437 layer_output = self.drop_path(layer_output) + hidden_states438 439 outputs = (layer_output,) + outputs440 441 return outputs442 443 444# Copied from transformers.models.beit.modeling_beit.BeitRelativePositionBias with Beit->Data2VecVision445class Data2VecVisionRelativePositionBias(nn.Module):446 def __init__(self, config: Data2VecVisionConfig, window_size: tuple) -> None:447 super().__init__()448 self.window_size = window_size449 self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3450 self.relative_position_bias_table = nn.Parameter(451 torch.zeros(self.num_relative_distance, config.num_attention_heads)452 ) # 2*Wh-1 * 2*Ww-1, nH453 # cls to token & token 2 cls & cls to cls454 455 # get pair-wise relative position index for each token inside the window456 coords_h = torch.arange(window_size[0])457 coords_w = torch.arange(window_size[1])458 coords = torch.stack(meshgrid([coords_h, coords_w], indexing="ij")) # 2, Wh, Ww459 coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww460 relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww461 relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2462 relative_coords[:, :, 0] += window_size[0] - 1 # shift to start from 0463 relative_coords[:, :, 1] += window_size[1] - 1464 relative_coords[:, :, 0] *= 2 * window_size[1] - 1465 relative_position_index = torch.zeros(466 size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype467 )468 relative_position_index[1:, 1:] = relative_coords.sum(-1) # Wh*Ww, Wh*Ww469 relative_position_index[0, 0:] = self.num_relative_distance - 3470 relative_position_index[0:, 0] = self.num_relative_distance - 2471 relative_position_index[0, 0] = self.num_relative_distance - 1472 473 self.register_buffer("relative_position_index", relative_position_index, persistent=False)474 475 def forward(self) -> torch.Tensor:476 relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(477 self.window_size[0] * self.window_size[1] + 1, self.window_size[0] * self.window_size[1] + 1, -1478 ) # Wh*Ww,Wh*Ww,nH479 480 return relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww481 482 483# Copied from transformers.models.beit.modeling_beit.BeitEncoder with Beit->Data2VecVision484class Data2VecVisionEncoder(nn.Module):485 def __init__(self, config: Data2VecVisionConfig, window_size: Optional[tuple] = None) -> None:486 super().__init__()487 self.config = config488 if config.use_shared_relative_position_bias:489 self.relative_position_bias = Data2VecVisionRelativePositionBias(config, window_size=window_size)490 else:491 self.relative_position_bias = None492 493 # stochastic depth decay rule494 dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers)]495 self.layer = nn.ModuleList(496 [497 Data2VecVisionLayer(498 config,499 window_size=window_size if config.use_relative_position_bias else None,500 drop_path_rate=dpr[i],501 )502 for i in range(config.num_hidden_layers)503 ]504 )505 self.gradient_checkpointing = False506 507 def forward(508 self,509 hidden_states: torch.Tensor,510 head_mask: Optional[torch.Tensor] = None,511 output_attentions: bool = False,512 output_hidden_states: bool = False,513 return_dict: bool = True,514 ) -> Union[tuple, BaseModelOutput]:515 all_hidden_states = () if output_hidden_states else None516 all_self_attentions = () if output_attentions else None517 518 for i, layer_module in enumerate(self.layer):519 if output_hidden_states:520 all_hidden_states = all_hidden_states + (hidden_states,)521 522 layer_head_mask = head_mask[i] if head_mask is not None else None523 524 if self.gradient_checkpointing and self.training:525 526 def create_custom_forward(module):527 def custom_forward(*inputs):528 return module(*inputs, output_attentions)529 530 return custom_forward531 532 layer_outputs = torch.utils.checkpoint.checkpoint(533 create_custom_forward(layer_module),534 hidden_states,535 layer_head_mask,536 )537 else:538 relative_position_bias = (539 self.relative_position_bias() if self.relative_position_bias is not None else None540 )541 layer_outputs = layer_module(hidden_states, layer_head_mask, output_attentions, relative_position_bias)542 543 hidden_states = layer_outputs[0]544 545 if output_attentions:546 all_self_attentions = all_self_attentions + (layer_outputs[1],)547 548 if output_hidden_states:549 all_hidden_states = all_hidden_states + (hidden_states,)550 551 if not return_dict:552 return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)553 return BaseModelOutput(554 last_hidden_state=hidden_states,555 hidden_states=all_hidden_states,556 attentions=all_self_attentions,557 )558 559 560# Copied from transformers.models.beit.modeling_beit.BeitPreTrainedModel with Beit->Data2VecVision,beit->data2vec_vision561class Data2VecVisionPreTrainedModel(PreTrainedModel):562 """563 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained564 models.565 """566 567 config_class = Data2VecVisionConfig568 base_model_prefix = "data2vec_vision"569 main_input_name = "pixel_values"570 supports_gradient_checkpointing = True571 572 def _init_weights(self, module):573 """Initialize the weights"""574 if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)):575 # Slightly different from the TF version which uses truncated_normal for initialization576 # cf https://github.com/pytorch/pytorch/pull/5617577 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)578 if module.bias is not None:579 module.bias.data.zero_()580 elif isinstance(module, nn.Embedding):581 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)582 if module.padding_idx is not None:583 module.weight.data[module.padding_idx].zero_()584 elif isinstance(module, nn.LayerNorm):585 module.bias.data.zero_()586 module.weight.data.fill_(1.0)587 588 def _set_gradient_checkpointing(self, module, value=False):589 if isinstance(module, Data2VecVisionEncoder):590 module.gradient_checkpointing = value591 592 593DATA2VEC_VISION_START_DOCSTRING = r"""594 This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it595 as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and596 behavior.597 598 Parameters:599 config ([`Data2VecVisionConfig`]): 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 604DATA2VEC_VISION_INPUTS_DOCSTRING = r"""605 Args:606 pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):607 Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See608 [`BeitImageProcessor.__call__`] for details.609 610 head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):611 Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:612 613 - 1 indicates the head is **not masked**,614 - 0 indicates the head is **masked**.615 616 output_attentions (`bool`, *optional*):617 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned618 tensors for more detail.619 output_hidden_states (`bool`, *optional*):620 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for621 more detail.622 return_dict (`bool`, *optional*):623 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.624"""625 626 627@add_start_docstrings(628 "The bare Data2VecVision Model transformer outputting raw hidden-states without any specific head on top.",629 DATA2VEC_VISION_START_DOCSTRING,630)631# Copied from transformers.models.beit.modeling_beit.BeitModel with BEIT->DATA2VEC_VISION,Beit->Data2VecVision,True->False632class Data2VecVisionModel(Data2VecVisionPreTrainedModel):633 def __init__(self, config: Data2VecVisionConfig, add_pooling_layer: bool = False) -> None:634 super().__init__(config)635 self.config = config636 637 self.embeddings = Data2VecVisionEmbeddings(config)638 self.encoder = Data2VecVisionEncoder(config, window_size=self.embeddings.patch_embeddings.patch_shape)639 640 self.layernorm = (641 nn.Identity() if config.use_mean_pooling else nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)642 )643 self.pooler = Data2VecVisionPooler(config) if add_pooling_layer else None644 645 # Initialize weights and apply final processing646 self.post_init()647 648 def get_input_embeddings(self):649 return self.embeddings.patch_embeddings650 651 def _prune_heads(self, heads_to_prune):652 """653 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base654 class PreTrainedModel655 """656 for layer, heads in heads_to_prune.items():657 self.encoder.layer[layer].attention.prune_heads(heads)658 659 @add_start_docstrings_to_model_forward(DATA2VEC_VISION_INPUTS_DOCSTRING)660 @add_code_sample_docstrings(661 checkpoint=_CHECKPOINT_FOR_DOC,662 output_type=Data2VecVisionModelOutputWithPooling,663 config_class=_CONFIG_FOR_DOC,664 modality="vision",665 expected_output=_EXPECTED_OUTPUT_SHAPE,666 )667 def forward(668 self,669 pixel_values: Optional[torch.Tensor] = None,670 bool_masked_pos: Optional[torch.BoolTensor] = None,671 head_mask: Optional[torch.Tensor] = None,672 output_attentions: Optional[bool] = None,673 output_hidden_states: Optional[bool] = None,674 return_dict: Optional[bool] = None,675 ) -> Union[tuple, Data2VecVisionModelOutputWithPooling]:676 r"""677 bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*):678 Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).679 """680 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions681 output_hidden_states = (682 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states683 )684 return_dict = return_dict if return_dict is not None else self.config.use_return_dict685 686 if pixel_values is None:687 raise ValueError("You have to specify pixel_values")688 689 # Prepare head mask if needed690 # 1.0 in head_mask indicate we keep the head691 # attention_probs has shape bsz x n_heads x N x N692 # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]693 # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]694 head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)695 696 embedding_output = self.embeddings(pixel_values, bool_masked_pos)697 698 encoder_outputs = self.encoder(699 embedding_output,700 head_mask=head_mask,701 output_attentions=output_attentions,702 output_hidden_states=output_hidden_states,703 return_dict=return_dict,704 )705 sequence_output = encoder_outputs[0]706 sequence_output = self.layernorm(sequence_output)707 pooled_output = self.pooler(sequence_output) if self.pooler is not None else None708 709 if not return_dict:710 head_outputs = (sequence_output, pooled_output) if pooled_output is not None else (sequence_output,)711 return head_outputs + encoder_outputs[1:]712 713 return Data2VecVisionModelOutputWithPooling(714 last_hidden_state=sequence_output,715 pooler_output=pooled_output,716 hidden_states=encoder_outputs.hidden_states,717 attentions=encoder_outputs.attentions,718 )719 720 721# Copied from transformers.models.beit.modeling_beit.BeitPooler with Beit->Data2VecVision722class Data2VecVisionPooler(nn.Module):723 def __init__(self, config: Data2VecVisionConfig) -> None:724 super().__init__()725 self.layernorm = (726 nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) if config.use_mean_pooling else None727 )728 729 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:730 if self.layernorm is not None:731 # Mean pool the final hidden states of the patch tokens732 patch_tokens = hidden_states[:, 1:, :]733 pooled_output = self.layernorm(patch_tokens.mean(1))734 else:735 # Pool by simply taking the final hidden state of the [CLS] token736 pooled_output = hidden_states[:, 0]737 738 return pooled_output739 740 741@add_start_docstrings(742 """743 Data2VecVision Model transformer with an image classification head on top (a linear layer on top of the average of744 the final hidden states of the patch tokens) e.g. for ImageNet.745 """,746 DATA2VEC_VISION_START_DOCSTRING,747)748# Copied from transformers.models.beit.modeling_beit.BeitForImageClassification with BEIT->DATA2VEC_VISION,Beit->Data2VecVision,beit->data2vec_vision749class Data2VecVisionForImageClassification(Data2VecVisionPreTrainedModel):750 def __init__(self, config: Data2VecVisionConfig) -> None:751 super().__init__(config)752 753 self.num_labels = config.num_labels754 self.data2vec_vision = Data2VecVisionModel(config, add_pooling_layer=True)755 756 # Classifier head757 self.classifier = nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()758 759 # Initialize weights and apply final processing760 self.post_init()761 762 @add_start_docstrings_to_model_forward(DATA2VEC_VISION_INPUTS_DOCSTRING)763 @add_code_sample_docstrings(764 checkpoint=_IMAGE_CLASS_CHECKPOINT,765 output_type=ImageClassifierOutput,766 config_class=_CONFIG_FOR_DOC,767 expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,768 )769 def forward(770 self,771 pixel_values: Optional[torch.Tensor] = None,772 head_mask: Optional[torch.Tensor] = None,773 labels: Optional[torch.Tensor] = None,774 output_attentions: Optional[bool] = None,775 output_hidden_states: Optional[bool] = None,776 return_dict: Optional[bool] = None,777 ) -> Union[tuple, ImageClassifierOutput]:778 r"""779 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):780 Labels for computing the image classification/regression loss. Indices should be in `[0, ...,781 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If782 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).783 """784 return_dict = return_dict if return_dict is not None else self.config.use_return_dict785 outputs = self.data2vec_vision(786 pixel_values,787 head_mask=head_mask,788 output_attentions=output_attentions,789 output_hidden_states=output_hidden_states,790 return_dict=return_dict,791 )792 793 pooled_output = outputs.pooler_output if return_dict else outputs[1]794 795 logits = self.classifier(pooled_output)796 797 loss = None798 if labels is not None:799 if self.config.problem_type is None:800 if self.num_labels == 1:801 self.config.problem_type = "regression"802 elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):803 self.config.problem_type = "single_label_classification"804 else:805 self.config.problem_type = "multi_label_classification"806 807 if self.config.problem_type == "regression":808 loss_fct = MSELoss()809 if self.num_labels == 1:810 loss = loss_fct(logits.squeeze(), labels.squeeze())811 else:812 loss = loss_fct(logits, labels)813 elif self.config.problem_type == "single_label_classification":814 loss_fct = CrossEntropyLoss()815 loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))816 elif self.config.problem_type == "multi_label_classification":817 loss_fct = BCEWithLogitsLoss()818 loss = loss_fct(logits, labels)819 if not return_dict:820 output = (logits,) + outputs[2:]821 return ((loss,) + output) if loss is not None else output822 823 return ImageClassifierOutput(824 loss=loss,825 logits=logits,826 hidden_states=outputs.hidden_states,827 attentions=outputs.attentions,828 )829 830 831# Copied from transformers.models.beit.modeling_beit.BeitConvModule with Beit->Data2VecVision832class Data2VecVisionConvModule(nn.Module):833 """834 A convolutional block that bundles conv/norm/activation layers. This block simplifies the usage of convolution835 layers, which are commonly used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU).836 837 Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.838 """839 840 def __init__(841 self,842 in_channels: int,843 out_channels: int,844 kernel_size: Union[int, Tuple[int, int]],845 padding: Union[int, Tuple[int, int], str] = 0,846 bias: bool = False,847 dilation: Union[int, Tuple[int, int]] = 1,848 ) -> None:849 super().__init__()850 self.conv = nn.Conv2d(851 in_channels=in_channels,852 out_channels=out_channels,853 kernel_size=kernel_size,854 padding=padding,855 bias=bias,856 dilation=dilation,857 )858 self.bn = nn.BatchNorm2d(out_channels)859 self.activation = nn.ReLU()860 861 def forward(self, input: torch.Tensor) -> torch.Tensor:862 output = self.conv(input)863 output = self.bn(output)864 output = self.activation(output)865 866 return output867 868 869# Copied from transformers.models.beit.modeling_beit.BeitPyramidPoolingBlock with Beit->Data2VecVision870class Data2VecVisionPyramidPoolingBlock(nn.Module):871 def __init__(self, pool_scale: int, in_channels: int, channels: int) -> None:872 super().__init__()873 self.layers = [874 nn.AdaptiveAvgPool2d(pool_scale),875 Data2VecVisionConvModule(in_channels, channels, kernel_size=1),876 ]877 for i, layer in enumerate(self.layers):878 self.add_module(str(i), layer)879 880 def forward(self, input: torch.Tensor) -> torch.Tensor:881 hidden_state = input882 for layer in self.layers:883 hidden_state = layer(hidden_state)884 return hidden_state885 886 887# Copied from transformers.models.beit.modeling_beit.BeitPyramidPoolingModule with Beit->Data2VecVision888class Data2VecVisionPyramidPoolingModule(nn.Module):889 """890 Pyramid Pooling Module (PPM) used in PSPNet.891 892 Args:893 pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid894 Module.895 in_channels (int): Input channels.896 channels (int): Channels after modules, before conv_seg.897 align_corners (bool): align_corners argument of F.interpolate.898 899 Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.900 """901 902 def __init__(self, pool_scales: Tuple[int, ...], in_channels: int, channels: int, align_corners: bool) -> None:903 super().__init__()904 self.pool_scales = pool_scales905 self.align_corners = align_corners906 self.in_channels = in_channels907 self.channels = channels908 self.blocks = []909 for i, pool_scale in enumerate(pool_scales):910 block = Data2VecVisionPyramidPoolingBlock(911 pool_scale=pool_scale, in_channels=in_channels, channels=channels912 )913 self.blocks.append(block)914 self.add_module(str(i), block)915 916 def forward(self, x: torch.Tensor) -> List[torch.Tensor]:917 ppm_outs = []918 for ppm in self.blocks:919 ppm_out = ppm(x)920 upsampled_ppm_out = nn.functional.interpolate(921 ppm_out, size=x.size()[2:], mode="bilinear", align_corners=self.align_corners922 )923 ppm_outs.append(upsampled_ppm_out)924 return ppm_outs925 926 927# Copied from transformers.models.beit.modeling_beit.BeitUperHead with Beit->Data2VecVision928class Data2VecVisionUperHead(nn.Module):929 """930 Unified Perceptual Parsing for Scene Understanding. This head is the implementation of931 [UPerNet](https://arxiv.org/abs/1807.10221).932 933 Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.934 """935 936 def __init__(self, config: Data2VecVisionConfig) -> None:937 super().__init__()938 939 self.pool_scales = config.pool_scales # e.g. (1, 2, 3, 6)940 self.in_channels = [config.hidden_size] * 4 # e.g. [768, 768, 768, 768]941 self.channels = config.hidden_size942 self.align_corners = False943 self.classifier = nn.Conv2d(self.channels, config.num_labels, kernel_size=1)944 945 # PSP Module946 self.psp_modules = Data2VecVisionPyramidPoolingModule(947 self.pool_scales,948 self.in_channels[-1],949 self.channels,950 align_corners=self.align_corners,951 )952 self.bottleneck = Data2VecVisionConvModule(953 self.in_channels[-1] + len(self.pool_scales) * self.channels,954 self.channels,955 kernel_size=3,956 padding=1,957 )958 # FPN Module959 self.lateral_convs = nn.ModuleList()960 self.fpn_convs = nn.ModuleList()961 for in_channels in self.in_channels[:-1]: # skip the top layer962 l_conv = Data2VecVisionConvModule(in_channels, self.channels, kernel_size=1)963 fpn_conv = Data2VecVisionConvModule(self.channels, self.channels, kernel_size=3, padding=1)964 self.lateral_convs.append(l_conv)965 self.fpn_convs.append(fpn_conv)966 967 self.fpn_bottleneck = Data2VecVisionConvModule(968 len(self.in_channels) * self.channels,969 self.channels,970 kernel_size=3,971 padding=1,972 )973 974 def psp_forward(self, inputs):975 x = inputs[-1]976 psp_outs = [x]977 psp_outs.extend(self.psp_modules(x))978 psp_outs = torch.cat(psp_outs, dim=1)979 output = self.bottleneck(psp_outs)980 981 return output982 983 def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:984 # build laterals985 laterals = [lateral_conv(encoder_hidden_states[i]) for i, lateral_conv in enumerate(self.lateral_convs)]986 987 laterals.append(self.psp_forward(encoder_hidden_states))988 989 # build top-down path990 used_backbone_levels = len(laterals)991 for i in range(used_backbone_levels - 1, 0, -1):992 prev_shape = laterals[i - 1].shape[2:]993 laterals[i - 1] = laterals[i - 1] + nn.functional.interpolate(994 laterals[i], size=prev_shape, mode="bilinear", align_corners=self.align_corners995 )996 997 # build outputs998 fpn_outs = [self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels - 1)]999 # append psp feature1000 fpn_outs.append(laterals[-1])1001 1002 for i in range(used_backbone_levels - 1, 0, -1):1003 fpn_outs[i] = nn.functional.interpolate(1004 fpn_outs[i], size=fpn_outs[0].shape[2:], mode="bilinear", align_corners=self.align_corners1005 )1006 fpn_outs = torch.cat(fpn_outs, dim=1)1007 output = self.fpn_bottleneck(fpn_outs)1008 output = self.classifier(output)1009 1010 return output1011 1012 1013# Copied from transformers.models.beit.modeling_beit.BeitFCNHead with Beit->Data2VecVision1014class Data2VecVisionFCNHead(nn.Module):1015 """1016 Fully Convolution Networks for Semantic Segmentation. This head is implemented of1017 [FCNNet](https://arxiv.org/abs/1411.4038>).1018 1019 Args:1020 config (Data2VecVisionConfig): Configuration.1021 in_channels1022 kernel_size (int): The kernel size for convs in the head. Default: 3.1023 dilation (int): The dilation rate for convs in the head. Default: 1.1024 1025 1026 Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.1027 """1028 1029 def __init__(1030 self,1031 config: Data2VecVisionConfig,1032 in_index: int = 2,1033 kernel_size: int = 3,1034 dilation: Union[int, Tuple[int, int]] = 1,1035 ) -> None:1036 super().__init__()1037 self.in_channels = config.hidden_size1038 self.channels = config.auxiliary_channels1039 self.num_convs = config.auxiliary_num_convs1040 self.concat_input = config.auxiliary_concat_input1041 self.in_index = in_index1042 1043 conv_padding = (kernel_size // 2) * dilation1044 convs = []1045 convs.append(1046 Data2VecVisionConvModule(1047 self.in_channels, self.channels, kernel_size=kernel_size, padding=conv_padding, dilation=dilation1048 )1049 )1050 for i in range(self.num_convs - 1):1051 convs.append(1052 Data2VecVisionConvModule(1053 self.channels, self.channels, kernel_size=kernel_size, padding=conv_padding, dilation=dilation1054 )1055 )1056 if self.num_convs == 0:1057 self.convs = nn.Identity()1058 else:1059 self.convs = nn.Sequential(*convs)1060 if self.concat_input:1061 self.conv_cat = Data2VecVisionConvModule(1062 self.in_channels + self.channels, self.channels, kernel_size=kernel_size, padding=kernel_size // 21063 )1064 1065 self.classifier = nn.Conv2d(self.channels, config.num_labels, kernel_size=1)1066 1067 def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:1068 # just take the relevant feature maps1069 hidden_states = encoder_hidden_states[self.in_index]1070 output = self.convs(hidden_states)1071 if self.concat_input:1072 output = self.conv_cat(torch.cat([hidden_states, output], dim=1))1073 output = self.classifier(output)1074 return output1075 1076 1077@add_start_docstrings(1078 """1079 Data2VecVision Model transformer with a semantic segmentation head on top e.g. for ADE20k, CityScapes.1080 """,1081 DATA2VEC_VISION_START_DOCSTRING,1082)1083# Copied from transformers.models.beit.modeling_beit.BeitForSemanticSegmentation with BEIT->DATA2VEC_VISION,Beit->Data2VecVision,microsoft/beit-base-finetuned-ade-640-640->facebook/data2vec-vision-base,beit->data2vec_vision1084class Data2VecVisionForSemanticSegmentation(Data2VecVisionPreTrainedModel):1085 def __init__(self, config: Data2VecVisionConfig) -> None:1086 super().__init__(config)1087 1088 self.num_labels = config.num_labels1089 self.data2vec_vision = Data2VecVisionModel(config, add_pooling_layer=False)1090 1091 # FPNs1092 self.fpn1 = nn.Sequential(1093 nn.ConvTranspose2d(config.hidden_size, config.hidden_size, kernel_size=2, stride=2),1094 nn.BatchNorm2d(config.hidden_size),1095 nn.GELU(),1096 nn.ConvTranspose2d(config.hidden_size, config.hidden_size, kernel_size=2, stride=2),1097 )1098 self.fpn2 = nn.Sequential(1099 nn.ConvTranspose2d(config.hidden_size, config.hidden_size, kernel_size=2, stride=2),1100 )1101 self.fpn3 = nn.Identity()1102 self.fpn4 = nn.MaxPool2d(kernel_size=2, stride=2)1103 1104 # Semantic segmentation head(s)1105 self.decode_head = Data2VecVisionUperHead(config)1106 self.auxiliary_head = Data2VecVisionFCNHead(config) if config.use_auxiliary_head else None1107 1108 # Initialize weights and apply final processing1109 self.post_init()1110 1111 def compute_loss(self, logits, auxiliary_logits, labels):1112 # upsample logits to the images' original size1113 upsampled_logits = nn.functional.interpolate(1114 logits, size=labels.shape[-2:], mode="bilinear", align_corners=False1115 )1116 if auxiliary_logits is not None:1117 upsampled_auxiliary_logits = nn.functional.interpolate(1118 auxiliary_logits, size=labels.shape[-2:], mode="bilinear", align_corners=False1119 )1120 # compute weighted loss1121 loss_fct = CrossEntropyLoss(ignore_index=self.config.semantic_loss_ignore_index)1122 main_loss = loss_fct(upsampled_logits, labels)1123 loss = main_loss1124 if auxiliary_logits is not None:1125 auxiliary_loss = loss_fct(upsampled_auxiliary_logits, labels)1126 loss += self.config.auxiliary_loss_weight * auxiliary_loss1127 1128 return loss1129 1130 @add_start_docstrings_to_model_forward(DATA2VEC_VISION_INPUTS_DOCSTRING)1131 @replace_return_docstrings(output_type=SemanticSegmenterOutput, config_class=_CONFIG_FOR_DOC)1132 def forward(1133 self,1134 pixel_values: Optional[torch.Tensor] = None,1135 head_mask: Optional[torch.Tensor] = None,1136 labels: Optional[torch.Tensor] = None,1137 output_attentions: Optional[bool] = None,1138 output_hidden_states: Optional[bool] = None,1139 return_dict: Optional[bool] = None,1140 ) -> Union[tuple, SemanticSegmenterOutput]:1141 r"""1142 labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*):1143 Ground truth semantic segmentation maps for computing the loss. Indices should be in `[0, ...,1144 config.num_labels - 1]`. If `config.num_labels > 1`, a classification loss is computed (Cross-Entropy).1145 1146 Returns:1147 1148 Examples:1149 1150 ```python1151 >>> from transformers import AutoImageProcessor, Data2VecVisionForSemanticSegmentation1152 >>> from PIL import Image1153 >>> import requests1154 1155 >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"1156 >>> image = Image.open(requests.get(url, stream=True).raw)1157 1158 >>> image_processor = AutoImageProcessor.from_pretrained("facebook/data2vec-vision-base")1159 >>> model = Data2VecVisionForSemanticSegmentation.from_pretrained("facebook/data2vec-vision-base")1160 1161 >>> inputs = image_processor(images=image, return_tensors="pt")1162 >>> outputs = model(**inputs)1163 >>> # logits are of shape (batch_size, num_labels, height, width)1164 >>> logits = outputs.logits1165 ```"""1166 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1167 output_hidden_states = (1168 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1169 )1170 1171 outputs = self.data2vec_vision(1172 pixel_values,1173 head_mask=head_mask,1174 output_attentions=output_attentions,1175 output_hidden_states=True, # we need the intermediate hidden states1176 return_dict=return_dict,1177 )1178 1179 encoder_hidden_states = outputs.hidden_states if return_dict else outputs[1]1180 1181 # only keep certain features, and reshape1182 # note that we do +1 as the encoder_hidden_states also includes the initial embeddings1183 features = [feature for idx, feature in enumerate(encoder_hidden_states) if idx + 1 in self.config.out_indices]1184 batch_size = pixel_values.shape[0]1185 patch_resolution = self.config.image_size // self.config.patch_size1186 features = [1187 x[:, 1:, :].permute(0, 2, 1).reshape(batch_size, -1, patch_resolution, patch_resolution) for x in features1188 ]1189 1190 # apply FPNs1191 ops = [self.fpn1, self.fpn2, self.fpn3, self.fpn4]1192 for i in range(len(features)):1193 features[i] = ops[i](features[i])1194 1195 logits = self.decode_head(features)1196 1197 auxiliary_logits = None1198 if self.auxiliary_head is not None:1199 auxiliary_logits = self.auxiliary_head(features)1200