Aluode/PerceptionLabPortable
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"""TF 2.0 Data2Vec Vision model."""16 17from __future__ import annotations18 19import collections.abc20import math21from dataclasses import dataclass22 23import numpy as np24import tensorflow as tf25 26from ...activations_tf import get_tf_activation27from ...modeling_tf_outputs import (28 TFBaseModelOutput,29 TFBaseModelOutputWithPooling,30 TFSemanticSegmenterOutput,31 TFSequenceClassifierOutput,32)33from ...modeling_tf_utils import (34 TFModelInputType,35 TFPreTrainedModel,36 TFSequenceClassificationLoss,37 get_initializer,38 keras,39 keras_serializable,40 unpack_inputs,41)42from ...tf_utils import shape_list, stable_softmax43from ...utils import (44 add_code_sample_docstrings,45 add_start_docstrings,46 add_start_docstrings_to_model_forward,47 logging,48 replace_return_docstrings,49)50from .configuration_data2vec_vision import Data2VecVisionConfig51 52 53logger = logging.get_logger(__name__)54 55# General docstring56_CONFIG_FOR_DOC = "Data2VecVisionConfig"57 58# Base docstring59_CHECKPOINT_FOR_DOC = "facebook/data2vec-vision-base"60_EXPECTED_OUTPUT_SHAPE = [1, 197, 768]61 62# Image classification docstring63_IMAGE_CLASS_CHECKPOINT = "facebook/data2vec-vision-base-ft1k"64_IMAGE_CLASS_EXPECTED_OUTPUT = "remote control, remote"65 66 67@dataclass68class TFData2VecVisionModelOutputWithPooling(TFBaseModelOutputWithPooling):69 """70 Class for outputs of [`TFData2VecVisionModel`].71 72 Args:73 last_hidden_state (`tf.Tensor` 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 (`tf.Tensor` 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(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):80 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape81 `(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(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):85 Tuple of `tf.Tensor` (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 last_hidden_state: tf.Tensor | None = None93 pooler_output: tf.Tensor | None = None94 hidden_states: tuple[tf.Tensor] | None = None95 attentions: tuple[tf.Tensor] | None = None96 97 98class TFData2VecVisionDropPath(keras.layers.Layer):99 """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).100 References:101 (1) github.com:rwightman/pytorch-image-models102 """103 104 def __init__(self, drop_path, **kwargs):105 super().__init__(**kwargs)106 self.drop_path = drop_path107 108 def call(self, x, training=None):109 if training:110 keep_prob = 1 - self.drop_path111 shape = (tf.shape(x)[0],) + (1,) * (len(tf.shape(x)) - 1)112 random_tensor = keep_prob + tf.random.uniform(shape, 0, 1)113 random_tensor = tf.floor(random_tensor)114 return (x / keep_prob) * random_tensor115 return x116 117 118class TFData2VecVisionEmbeddings(keras.layers.Layer):119 """120 Construct the CLS token, position and patch embeddings. Optionally, also the mask token.121 122 """123 124 def __init__(self, config: Data2VecVisionConfig, **kwargs):125 super().__init__(**kwargs)126 self.config = config127 128 self.patch_embeddings = TFData2VecVisionPatchEmbeddings(config, name="patch_embeddings")129 self.num_patches = self.patch_embeddings.num_patches130 self.config = config131 132 self.dropout = keras.layers.Dropout(config.hidden_dropout_prob)133 134 def build(self, input_shape=None):135 self.cls_token = self.add_weight(136 shape=(1, 1, self.config.hidden_size),137 initializer=tf.random_normal_initializer(stddev=self.config.initializer_range),138 trainable=True,139 name="cls_token",140 )141 if self.config.use_mask_token:142 self.mask_token = self.add_weight(143 shape=(1, 1, self.config.hidden_size),144 initializer=tf.random_normal_initializer(stddev=self.config.initializer_range),145 trainable=True,146 name="mask_token",147 )148 else:149 self.mask_token = None150 151 if self.config.use_absolute_position_embeddings:152 self.position_embeddings = self.add_weight(153 shape=(1, self.num_patches + 1, self.config.hidden_size),154 initializer=tf.random_normal_initializer(stddev=self.config.initializer_range),155 trainable=True,156 name="position_embeddings",157 )158 else:159 self.position_embeddings = None160 161 if self.built:162 return163 self.built = True164 if getattr(self, "patch_embeddings", None) is not None:165 with tf.name_scope(self.patch_embeddings.name):166 self.patch_embeddings.build(None)167 168 def call(self, pixel_values: tf.Tensor, bool_masked_pos: tf.Tensor | None = None) -> tf.Tensor:169 embeddings = self.patch_embeddings(pixel_values)170 batch_size, seq_len, projection_dim = shape_list(embeddings)171 172 cls_tokens = tf.tile(self.cls_token, (batch_size, 1, 1))173 174 if bool_masked_pos is not None:175 mask_tokens = tf.broadcast_to(self.mask_token, (batch_size, seq_len, projection_dim))176 # replace the masked visual tokens by mask_tokens177 w = bool_masked_pos[..., None]178 w = tf.cast(w, mask_tokens.dtype)179 # since TF doesn't support eager tensor assignment180 embeddings = embeddings * (1 - w) + mask_tokens * w181 182 embeddings = tf.concat([cls_tokens, embeddings], axis=1)183 if self.position_embeddings is not None:184 embeddings = embeddings + self.position_embeddings185 embeddings = self.dropout(embeddings)186 187 return embeddings188 189 190class TFData2VecVisionPatchEmbeddings(keras.layers.Layer):191 """192 Image to Patch Embedding.193 """194 195 def __init__(self, config: Data2VecVisionConfig, **kwargs):196 super().__init__(**kwargs)197 self.config = config198 199 image_size, patch_size = config.image_size, config.patch_size200 num_channels, hidden_size = config.num_channels, config.hidden_size201 202 image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)203 patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)204 num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])205 patch_shape = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])206 self.image_size = image_size207 self.patch_size = patch_size208 self.num_patches = num_patches209 self.patch_shape = patch_shape210 self.num_channels = num_channels211 212 self.projection = keras.layers.Conv2D(213 filters=hidden_size,214 kernel_size=patch_size,215 strides=patch_size,216 padding="valid",217 data_format="channels_last",218 kernel_initializer="glorot_uniform", # following torch.nn.Linear219 bias_initializer="zeros",220 name="projection",221 )222 223 def call(self, pixel_values: tf.Tensor, training: bool = False) -> tf.Tensor:224 batch_size, num_channels, height, width = shape_list(pixel_values)225 if tf.executing_eagerly():226 if num_channels != self.num_channels:227 raise ValueError(228 "Make sure that the channel dimension of the pixel values match with the one set in the"229 " configuration."230 )231 if height != self.image_size[0] or width != self.image_size[1]:232 raise ValueError(233 f"Input image size ({height}*{width}) doesn't match model"234 f" ({self.image_size[0]}*{self.image_size[1]})."235 )236 237 # When running on CPU, `keras.layers.Conv2D` doesn't support `NCHW` format.238 # So change the input format from `NCHW` to `NHWC`.239 # shape = (batch_size, in_height, in_width, in_channels=num_channels)240 pixel_values = tf.transpose(pixel_values, perm=(0, 2, 3, 1))241 242 projection = self.projection(pixel_values)243 244 # Change the 2D spatial dimensions to a single temporal dimension.245 # shape = (batch_size, num_patches, out_channels=embed_dim)246 num_patches = (width // self.patch_size[1]) * (height // self.patch_size[0])247 248 return tf.reshape(tensor=projection, shape=(batch_size, num_patches, -1))249 250 def build(self, input_shape=None):251 if self.built:252 return253 self.built = True254 if getattr(self, "projection", None) is not None:255 with tf.name_scope(self.projection.name):256 self.projection.build([None, None, None, self.num_channels])257 258 259class TFData2VecVisionSelfAttention(keras.layers.Layer):260 def __init__(self, config: Data2VecVisionConfig, window_size: tuple | None = None, **kwargs):261 super().__init__(**kwargs)262 263 if config.hidden_size % config.num_attention_heads != 0:264 raise ValueError(265 f"The hidden size ({config.hidden_size}) is not a multiple of the number "266 f"of attention heads ({config.num_attention_heads})"267 )268 269 self.num_attention_heads = config.num_attention_heads270 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)271 self.all_head_size = self.num_attention_heads * self.attention_head_size272 self.sqrt_att_head_size = math.sqrt(self.attention_head_size)273 274 self.query = keras.layers.Dense(275 units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query"276 )277 self.key = keras.layers.Dense(278 units=self.all_head_size,279 kernel_initializer=get_initializer(config.initializer_range),280 name="key",281 use_bias=False,282 )283 self.value = keras.layers.Dense(284 units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value"285 )286 self.dropout = keras.layers.Dropout(rate=config.attention_probs_dropout_prob)287 288 if window_size:289 self.relative_position_bias = TFData2VecVisionRelativePositionBias(290 config, window_size=window_size, name="relative_position_bias"291 )292 else:293 self.relative_position_bias = None294 self.config = config295 296 def transpose_for_scores(self, tensor: tf.Tensor, batch_size: int) -> tf.Tensor:297 # Reshape from [batch_size, seq_length, all_head_size] to [batch_size, seq_length, num_attention_heads, attention_head_size]298 tensor = tf.reshape(tensor=tensor, shape=(batch_size, -1, self.num_attention_heads, self.attention_head_size))299 300 # Transpose the tensor from [batch_size, seq_length, num_attention_heads, attention_head_size] to [batch_size, num_attention_heads, seq_length, attention_head_size]301 return tf.transpose(tensor, perm=[0, 2, 1, 3])302 303 def call(304 self,305 hidden_states: tf.Tensor,306 head_mask: tf.Tensor,307 output_attentions: bool,308 relative_position_bias: TFData2VecVisionRelativePositionBias | None = None,309 training: bool = False,310 ) -> tuple[tf.Tensor]:311 batch_size = shape_list(hidden_states)[0]312 mixed_query_layer = self.query(inputs=hidden_states)313 mixed_key_layer = self.key(inputs=hidden_states)314 mixed_value_layer = self.value(inputs=hidden_states)315 query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)316 key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)317 value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)318 319 # Take the dot product between "query" and "key" to get the raw attention scores.320 # (batch size, num_heads, seq_len_q, seq_len_k)321 attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)322 attention_scores = attention_scores / self.sqrt_att_head_size323 324 # Add relative position bias if present.325 if self.relative_position_bias is not None:326 # Passing `0.0` to the `relative_position_bias()` layer because otherwise Keras327 # might complain about `Layer.call()` not being invoked properly. In this case this input328 # i.e., 0.0 is not going to be used in any calculations so we're safe.329 attention_scores = attention_scores + self.relative_position_bias(0.0)[None, ...]330 331 # Add shared relative position bias if provided.332 if relative_position_bias is not None:333 attention_scores = attention_scores + relative_position_bias334 335 # Normalize the attention scores to probabilities.336 attention_probs = stable_softmax(logits=attention_scores, axis=-1)337 338 # This is actually dropping out entire tokens to attend to, which might339 # seem a bit unusual, but is taken from the original Transformer paper.340 attention_probs = self.dropout(inputs=attention_probs, training=training)341 342 # Mask heads if we want to343 if head_mask is not None:344 attention_probs = tf.multiply(attention_probs, head_mask)345 346 attention_output = tf.matmul(attention_probs, value_layer)347 attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])348 349 # (batch_size, seq_len_q, all_head_size)350 attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.all_head_size))351 outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)352 353 return outputs354 355 def build(self, input_shape=None):356 if self.built:357 return358 self.built = True359 if getattr(self, "query", None) is not None:360 with tf.name_scope(self.query.name):361 self.query.build([None, None, self.config.hidden_size])362 if getattr(self, "key", None) is not None:363 with tf.name_scope(self.key.name):364 self.key.build([None, None, self.config.hidden_size])365 if getattr(self, "value", None) is not None:366 with tf.name_scope(self.value.name):367 self.value.build([None, None, self.config.hidden_size])368 if getattr(self, "relative_position_bias", None) is not None:369 with tf.name_scope(self.relative_position_bias.name):370 self.relative_position_bias.build(None)371 372 373class TFData2VecVisionSelfOutput(keras.layers.Layer):374 """375 The residual connection is defined in TFData2VecVisionLayer instead of here (as is the case with other models), due376 to the layernorm applied before each block.377 """378 379 def __init__(self, config: Data2VecVisionConfig, **kwargs):380 super().__init__(**kwargs)381 382 self.dense = keras.layers.Dense(383 units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"384 )385 self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)386 self.config = config387 388 def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, gamma=None, training: bool = False) -> tf.Tensor:389 hidden_states = self.dense(inputs=hidden_states)390 hidden_states = self.dropout(inputs=hidden_states, training=training)391 392 return hidden_states393 394 def build(self, input_shape=None):395 if self.built:396 return397 self.built = True398 if getattr(self, "dense", None) is not None:399 with tf.name_scope(self.dense.name):400 self.dense.build([None, None, self.config.hidden_size])401 402 403class TFData2VecVisionAttention(keras.layers.Layer):404 def __init__(self, config: Data2VecVisionConfig, window_size: tuple | None = None, **kwargs):405 super().__init__(**kwargs)406 407 self.attention = TFData2VecVisionSelfAttention(config, window_size=window_size, name="attention")408 self.dense_output = TFData2VecVisionSelfOutput(config, name="output")409 410 def prune_heads(self, heads):411 raise NotImplementedError412 413 def call(414 self,415 input_tensor: tf.Tensor,416 head_mask: tf.Tensor,417 output_attentions: bool,418 relative_position_bias: TFData2VecVisionRelativePositionBias | None = None,419 training: bool = False,420 ) -> tuple[tf.Tensor]:421 self_outputs = self.attention(422 hidden_states=input_tensor,423 head_mask=head_mask,424 output_attentions=output_attentions,425 relative_position_bias=relative_position_bias,426 training=training,427 )428 attention_output = self.dense_output(429 hidden_states=self_outputs[0], input_tensor=input_tensor, training=training430 )431 outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them432 433 return outputs434 435 def build(self, input_shape=None):436 if self.built:437 return438 self.built = True439 if getattr(self, "attention", None) is not None:440 with tf.name_scope(self.attention.name):441 self.attention.build(None)442 if getattr(self, "dense_output", None) is not None:443 with tf.name_scope(self.dense_output.name):444 self.dense_output.build(None)445 446 447# Copied from transformers.models.vit.modeling_tf_vit.TFViTIntermediate with ViT->Data2VecVision448class TFData2VecVisionIntermediate(keras.layers.Layer):449 def __init__(self, config: Data2VecVisionConfig, **kwargs):450 super().__init__(**kwargs)451 452 self.dense = keras.layers.Dense(453 units=config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"454 )455 456 if isinstance(config.hidden_act, str):457 self.intermediate_act_fn = get_tf_activation(config.hidden_act)458 else:459 self.intermediate_act_fn = config.hidden_act460 self.config = config461 462 def call(self, hidden_states: tf.Tensor) -> tf.Tensor:463 hidden_states = self.dense(inputs=hidden_states)464 hidden_states = self.intermediate_act_fn(hidden_states)465 466 return hidden_states467 468 def build(self, input_shape=None):469 if self.built:470 return471 self.built = True472 if getattr(self, "dense", None) is not None:473 with tf.name_scope(self.dense.name):474 self.dense.build([None, None, self.config.hidden_size])475 476 477class TFData2VecVisionOutput(keras.layers.Layer):478 def __init__(self, config: Data2VecVisionConfig, **kwargs):479 super().__init__(**kwargs)480 481 self.dense = keras.layers.Dense(482 units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"483 )484 self.dropout = keras.layers.Dropout(rate=config.hidden_dropout_prob)485 self.config = config486 487 def call(self, hidden_states: tf.Tensor, training: bool = False) -> tf.Tensor:488 hidden_states = self.dense(inputs=hidden_states)489 hidden_states = self.dropout(inputs=hidden_states, training=training)490 491 return hidden_states492 493 def build(self, input_shape=None):494 if self.built:495 return496 self.built = True497 if getattr(self, "dense", None) is not None:498 with tf.name_scope(self.dense.name):499 self.dense.build([None, None, self.config.intermediate_size])500 501 502class TFData2VecVisionLayer(keras.layers.Layer):503 """This corresponds to the Block class in the timm implementation."""504 505 def __init__(506 self, config: Data2VecVisionConfig, window_size: tuple | None = None, drop_path_rate: float = 0.0, **kwargs507 ):508 super().__init__(**kwargs)509 self.config = config510 511 self.attention = TFData2VecVisionAttention(config, window_size=window_size, name="attention")512 self.intermediate = TFData2VecVisionIntermediate(config, name="intermediate")513 self.data2vec_output = TFData2VecVisionOutput(config, name="output")514 515 self.layernorm_before = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm_before")516 self.layernorm_after = keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm_after")517 # Using `layers.Activation` instead of `tf.identity` to better control `training`518 # behaviour.519 self.drop_path = (520 TFData2VecVisionDropPath(drop_path_rate, name="drop_path")521 if drop_path_rate > 0.0522 else keras.layers.Activation("linear", name="drop_path")523 )524 self.init_values = config.layer_scale_init_value525 526 def build(self, input_shape: tf.TensorShape = None):527 if self.init_values > 0:528 self.lambda_1 = self.add_weight(529 shape=(self.config.hidden_size),530 initializer="ones",531 trainable=True,532 name="lambda_1",533 )534 self.lambda_2 = self.add_weight(535 shape=(self.config.hidden_size),536 initializer="ones",537 trainable=True,538 name="lambda_2",539 )540 self.lambda_1.assign(self.init_values * tf.ones(self.config.hidden_size))541 self.lambda_2.assign(self.init_values * tf.ones(self.config.hidden_size))542 else:543 self.lambda_1, self.lambda_2 = None, None544 545 if self.built:546 return547 self.built = True548 if getattr(self, "attention", None) is not None:549 with tf.name_scope(self.attention.name):550 self.attention.build(None)551 if getattr(self, "intermediate", None) is not None:552 with tf.name_scope(self.intermediate.name):553 self.intermediate.build(None)554 if getattr(self, "data2vec_output", None) is not None:555 with tf.name_scope(self.data2vec_output.name):556 self.data2vec_output.build(None)557 if getattr(self, "layernorm_before", None) is not None:558 with tf.name_scope(self.layernorm_before.name):559 self.layernorm_before.build([None, None, self.config.hidden_size])560 if getattr(self, "layernorm_after", None) is not None:561 with tf.name_scope(self.layernorm_after.name):562 self.layernorm_after.build([None, None, self.config.hidden_size])563 if getattr(self, "drop_path", None) is not None:564 with tf.name_scope(self.drop_path.name):565 self.drop_path.build(None)566 567 def call(568 self,569 hidden_states: tf.Tensor,570 head_mask: tf.Tensor,571 output_attentions: bool,572 relative_position_bias: TFData2VecVisionRelativePositionBias | None = None,573 training: bool = False,574 ) -> tuple[tf.Tensor]:575 self_attention_outputs = self.attention(576 # in Data2VecVision, layernorm is applied before self-attention577 input_tensor=self.layernorm_before(inputs=hidden_states),578 head_mask=head_mask,579 output_attentions=output_attentions,580 relative_position_bias=relative_position_bias,581 training=training,582 )583 attention_output = self_attention_outputs[0]584 outputs = self_attention_outputs[1:] # add self attentions if we output attention weights585 586 # apply lambda_1 if present587 if self.lambda_1 is not None:588 attention_output = self.lambda_1 * attention_output589 590 # first residual connection591 hidden_states = self.drop_path(attention_output) + hidden_states592 593 # in Data2VecVision, layernorm is also applied after self-attention594 layer_output = self.layernorm_after(hidden_states)595 596 layer_output = self.intermediate(layer_output)597 layer_output = self.data2vec_output(layer_output)598 599 if self.lambda_2 is not None:600 layer_output = self.lambda_2 * layer_output601 602 # second residual connection603 layer_output = self.drop_path(layer_output) + hidden_states604 605 outputs = (layer_output,) + outputs606 607 return outputs608 609 610# Taken and modified from here:611# https://github.com/leondgarse/keras_cv_attention_models/blob/main/keras_cv_attention_models/beit/beit.py#L28612class TFData2VecVisionRelativePositionBias(keras.layers.Layer):613 def __init__(self, config: Data2VecVisionConfig, window_size: tuple, **kwargs) -> None:614 super().__init__(**kwargs)615 self.config = config616 617 self.window_size = window_size618 # +3 for cls_token_pos_len619 # window_size can be something like (14, 14)620 self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3621 622 self.relative_position_index = self.get_position_index()623 624 def build(self, input_shape):625 self.relative_position_bias_table = self.add_weight(626 shape=(self.num_relative_distance, self.config.num_attention_heads),627 initializer="zeros",628 trainable=True,629 name="relative_position_bias_table",630 ) # [2*Wh-1 * 2*Ww-1, nH]631 # cls to token & token 2 cls & cls to cls632 633 super().build(input_shape)634 635 def get_position_index(self):636 # get pair-wise relative position index for each token inside the window637 xx, yy = tf.meshgrid(range(self.window_size[0]), range(self.window_size[1]))638 coords = tf.stack([yy, xx], axis=0) # [2, Wh, Ww]639 coords_flatten = tf.reshape(coords, [2, -1]) # [2, Wh*Ww]640 641 relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # [2, Wh*Ww, Wh*Ww]642 relative_coords = tf.transpose(relative_coords, perm=[1, 2, 0]) # [Wh*Ww, Wh*Ww, 2]643 644 xx = (relative_coords[:, :, 0] + self.window_size[0] - 1) * (2 * self.window_size[1] - 1)645 yy = relative_coords[:, :, 1] + self.window_size[1] - 1646 relative_coords = tf.stack([xx, yy], axis=-1)647 648 relative_position_index = tf.reduce_sum(relative_coords, axis=-1) # [Wh*Ww, Wh*Ww]649 650 top = tf.ones((1, relative_position_index.shape[1]), dtype=relative_position_index.dtype) * (651 self.num_relative_distance - 3652 )653 left = tf.ones((relative_position_index.shape[0], 1), dtype=relative_position_index.dtype) * (654 self.num_relative_distance - 2655 )656 corner = tf.ones((1, 1), dtype=relative_position_index.dtype) * (self.num_relative_distance - 1)657 658 left_corner = tf.concat([corner, left], axis=0)659 relative_position_index = tf.concat([top, relative_position_index], axis=0)660 relative_position_index = tf.concat([left_corner, relative_position_index], axis=1) # [Wh*Ww + 1, Wh*Ww + 1]661 return relative_position_index662 663 def call(self, inputs=None) -> tf.Tensor:664 relative_position_bias = tf.gather(self.relative_position_bias_table, self.relative_position_index, axis=0)665 return tf.transpose(relative_position_bias, [2, 0, 1])666 667 668class TFData2VecVisionEncoder(keras.layers.Layer):669 def __init__(self, config: Data2VecVisionConfig, window_size: tuple | None = None, **kwargs):670 super().__init__(**kwargs)671 self.config = config672 if config.use_shared_relative_position_bias:673 self.relative_position_bias = TFData2VecVisionRelativePositionBias(674 config, window_size=window_size, name="relative_position_bias"675 )676 else:677 self.relative_position_bias = None678 679 # stochastic depth decay rule680 dpr = list(tf.linspace(0.0, config.drop_path_rate, config.num_hidden_layers))681 self.layer = [682 TFData2VecVisionLayer(683 config,684 window_size=window_size if config.use_relative_position_bias else None,685 drop_path_rate=dpr[i],686 name=f"layer_._{i}",687 )688 for i in range(config.num_hidden_layers)689 ]690 691 def call(692 self,693 hidden_states: tf.Tensor,694 head_mask: tf.Tensor | None = None,695 output_attentions: bool = False,696 output_hidden_states: bool = False,697 return_dict: bool = True,698 ) -> tuple | TFBaseModelOutput:699 all_hidden_states = () if output_hidden_states else None700 all_self_attentions = () if output_attentions else None701 702 for i, layer_module in enumerate(self.layer):703 if output_hidden_states:704 all_hidden_states = all_hidden_states + (hidden_states,)705 706 layer_head_mask = head_mask[i] if head_mask is not None else None707 # Passing `0.0` to the `relative_position_bias()` layer because otherwise Keras708 # might complain about `Layer.call()` not being invoked properly. In this case this input709 # i.e., 0.0 is not going to be used in any calculations so we're safe.710 relative_position_bias = (711 self.relative_position_bias(0.0) if self.relative_position_bias is not None else None712 )713 layer_outputs = layer_module(hidden_states, layer_head_mask, output_attentions, relative_position_bias)714 715 hidden_states = layer_outputs[0]716 717 if output_attentions:718 all_self_attentions = all_self_attentions + (layer_outputs[1],)719 720 if output_hidden_states:721 all_hidden_states = all_hidden_states + (hidden_states,)722 723 if not return_dict:724 return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)725 726 return TFBaseModelOutput(727 last_hidden_state=hidden_states,728 hidden_states=all_hidden_states,729 attentions=all_self_attentions,730 )731 732 def build(self, input_shape=None):733 if self.built:734 return735 self.built = True736 if getattr(self, "relative_position_bias", None) is not None:737 with tf.name_scope(self.relative_position_bias.name):738 self.relative_position_bias.build(None)739 if getattr(self, "layer", None) is not None:740 for layer in self.layer:741 with tf.name_scope(layer.name):742 layer.build(None)743 744 745@keras_serializable746class TFData2VecVisionMainLayer(keras.layers.Layer):747 config_class = Data2VecVisionConfig748 749 def __init__(self, config: Data2VecVisionConfig, add_pooling_layer: bool = True, **kwargs):750 super().__init__(**kwargs)751 752 self.config = config753 self.add_pooling_layer = add_pooling_layer754 755 self.embeddings = TFData2VecVisionEmbeddings(config, name="embeddings")756 self.encoder = TFData2VecVisionEncoder(757 config, window_size=self.embeddings.patch_embeddings.patch_shape, name="encoder"758 )759 self.layernorm = (760 tf.identity761 if config.use_mean_pooling762 else keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")763 )764 765 # We are setting the `data_format` like so because from here on we will revert to the766 # NCHW output format767 self.pooler = TFData2VecVisionPooler(config, name="pooler") if add_pooling_layer else None768 769 def get_input_embeddings(self) -> keras.layers.Layer:770 return self.embeddings.patch_embeddings771 772 def _prune_heads(self, heads_to_prune):773 """774 Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base775 class PreTrainedModel776 """777 raise NotImplementedError778 779 @unpack_inputs780 def call(781 self,782 pixel_values: tf.Tensor | None = None,783 bool_masked_pos: tf.Tensor | None = None,784 head_mask: tf.Tensor | None = None,785 output_attentions: bool | None = None,786 output_hidden_states: bool | None = None,787 return_dict: bool | None = None,788 training: bool = False,789 ) -> tuple | TFData2VecVisionModelOutputWithPooling:790 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions791 output_hidden_states = (792 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states793 )794 return_dict = return_dict if return_dict is not None else self.config.use_return_dict795 796 if pixel_values is None:797 raise ValueError("You have to specify pixel_values")798 799 # Prepare head mask if needed800 # 1.0 in head_mask indicate we keep the head801 # attention_probs has shape bsz x n_heads x N x N802 # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]803 # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]804 if head_mask is not None:805 raise NotImplementedError806 else:807 head_mask = [None] * self.config.num_hidden_layers808 809 embedding_output = self.embeddings(pixel_values, bool_masked_pos, training=training)810 811 encoder_outputs = self.encoder(812 embedding_output,813 head_mask=head_mask,814 output_attentions=output_attentions,815 output_hidden_states=output_hidden_states,816 return_dict=return_dict,817 training=training,818 )819 820 sequence_output = encoder_outputs[0]821 sequence_output = self.layernorm(sequence_output)822 pooled_output = self.pooler(sequence_output) if self.pooler is not None else None823 824 if not return_dict:825 head_outputs = (sequence_output, pooled_output) if pooled_output is not None else (sequence_output,)826 return head_outputs + encoder_outputs[1:]827 828 return TFData2VecVisionModelOutputWithPooling(829 last_hidden_state=sequence_output,830 pooler_output=pooled_output,831 hidden_states=encoder_outputs.hidden_states,832 attentions=encoder_outputs.attentions,833 )834 835 def build(self, input_shape=None):836 if self.built:837 return838 self.built = True839 if getattr(self, "embeddings", None) is not None:840 with tf.name_scope(self.embeddings.name):841 self.embeddings.build(None)842 if getattr(self, "encoder", None) is not None:843 with tf.name_scope(self.encoder.name):844 self.encoder.build(None)845 if getattr(self, "layernorm", None) is not None:846 if hasattr(self.layernorm, "name"):847 with tf.name_scope(self.layernorm.name):848 self.layernorm.build((None, self.config.hidden_size))849 if getattr(self, "pooler", None) is not None:850 with tf.name_scope(self.pooler.name):851 self.pooler.build(None)852 853 854class TFData2VecVisionPooler(keras.layers.Layer):855 def __init__(self, config: Data2VecVisionConfig, **kwargs):856 super().__init__(**kwargs)857 self.layernorm = (858 keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")859 if config.use_mean_pooling860 else None861 )862 self.config = config863 864 def call(self, hidden_states: tf.Tensor) -> tf.Tensor:865 if self.layernorm is not None:866 # Mean pool the final hidden states of the patch tokens867 patch_tokens = hidden_states[:, 1:, :]868 pooled_output = self.layernorm(tf.reduce_mean(patch_tokens, axis=1))869 else:870 # Pool by simply taking the final hidden state of the [CLS] token871 pooled_output = hidden_states[:, 0]872 873 return pooled_output874 875 def build(self, input_shape=None):876 if self.built:877 return878 self.built = True879 if getattr(self, "layernorm", None) is not None:880 if hasattr(self.layernorm, "name"):881 with tf.name_scope(self.layernorm.name):882 self.layernorm.build((None, self.config.hidden_size))883 884 885class TFData2VecVisionPreTrainedModel(TFPreTrainedModel):886 """887 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained888 models.889 """890 891 config_class = Data2VecVisionConfig892 base_model_prefix = "data2vec_vision"893 main_input_name = "pixel_values"894 _keys_to_ignore_on_load_unexpected = [r"relative_position_index"]895 896 897DATA2VEC_VISION_START_DOCSTRING = r"""898 This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the generic methods the899 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads900 etc.).901 902 This model is also a [keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use it903 as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage and904 behavior.905 906 <Tip>907 908 TensorFlow models and layers in `transformers` accept two formats as input:909 910 - having all inputs as keyword arguments (like PyTorch models), or911 - having all inputs as a list, tuple or dict in the first positional argument.912 913 The reason the second format is supported is that Keras methods prefer this format when passing inputs to models914 and layers. Because of this support, when using methods like `model.fit()` things should "just work" for you - just915 pass your inputs and labels in any format that `model.fit()` supports! If, however, you want to use the second916 format outside of Keras methods like `fit()` and `predict()`, such as when creating your own layers or models with917 the Keras `Functional` API, there are three possibilities you can use to gather all the input Tensors in the first918 positional argument:919 920 - a single Tensor with `pixel_values` only and nothing else: `model(pixel_values)`921 - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:922 `model([pixel_values, attention_mask])` or `model([pixel_values, attention_mask, token_type_ids])`923 - a dictionary with one or several input Tensors associated to the input names given in the docstring:924 `model({"pixel_values": pixel_values, "token_type_ids": token_type_ids})`925 926 Note that when creating models and layers with927 [subclassing](https://keras.io/guides/making_new_layers_and_models_via_subclassing/) then you don't need to worry928 about any of this, as you can just pass inputs like you would to any other Python function!929 930 </Tip>931 932 Args:933 config ([`Data2VecVisionConfig`]): Model configuration class with all the parameters of the model.934 Initializing with a config file does not load the weights associated with the model, only the935 configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.936"""937 938DATA2VEC_VISION_INPUTS_DOCSTRING = r"""939 Args:940 pixel_values (`np.ndarray`, `tf.Tensor`, `list[tf.Tensor]` `dict[str, tf.Tensor]` or `dict[str, np.ndarray]` and each example must have the shape `(batch_size, num_channels, height, width)`):941 Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See942 [`BeitImageProcessor.__call__`] for details.943 944 head_mask (`np.ndarray` or `tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):945 Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:946 - 1 indicates the head is **not masked**,947 - 0 indicates the head is **masked**.948 949 output_attentions (`bool`, *optional*):950 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned951 tensors for more detail.952 953 output_hidden_states (`bool`, *optional*):954 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for955 more detail.956 957 return_dict (`bool`, *optional*):958 Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. This argument can be used959 in eager mode, in graph mode the value will always be set to True.960 961 training (`bool`, *optional*, defaults to `False``):962 Whether or not to use the model in training mode (some modules like dropout modules have different963 behaviors between training and evaluation).964"""965 966 967@add_start_docstrings(968 "The bare Data2VecVision Model transformer outputting raw hidden-states without any specific head on top.",969 DATA2VEC_VISION_START_DOCSTRING,970)971class TFData2VecVisionModel(TFData2VecVisionPreTrainedModel):972 def __init__(self, config: Data2VecVisionConfig, add_pooling_layer: bool = False, *inputs, **kwargs):973 super().__init__(config, *inputs, **kwargs)974 self.config = config975 976 self.data2vec_vision = TFData2VecVisionMainLayer(977 config, add_pooling_layer=add_pooling_layer, name="data2vec_vision"978 )979 980 def get_input_embeddings(self):981 return self.data2vec_vision.get_input_embeddings()982 983 @unpack_inputs984 @add_start_docstrings_to_model_forward(DATA2VEC_VISION_INPUTS_DOCSTRING)985 @add_code_sample_docstrings(986 checkpoint=_CHECKPOINT_FOR_DOC,987 output_type=TFData2VecVisionModelOutputWithPooling,988 config_class=_CONFIG_FOR_DOC,989 modality="vision",990 expected_output=_EXPECTED_OUTPUT_SHAPE,991 )992 def call(993 self,994 pixel_values: TFModelInputType | None = None,995 bool_masked_pos: tf.Tensor | None = None,996 head_mask: np.ndarray | tf.Tensor | None = None,997 output_attentions: bool | None = None,998 output_hidden_states: bool | None = None,999 return_dict: bool | None = None,1000 training: bool = False,1001 ) -> tuple | TFData2VecVisionModelOutputWithPooling:1002 r"""1003 bool_masked_pos (`tf.Tensor` of shape `(batch_size, num_patches)`, *optional*):1004 Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).1005 """1006 outputs = self.data2vec_vision(1007 pixel_values=pixel_values,1008 bool_masked_pos=bool_masked_pos,1009 head_mask=head_mask,1010 output_attentions=output_attentions,1011 output_hidden_states=output_hidden_states,1012 return_dict=return_dict,1013 training=training,1014 )1015 1016 return outputs1017 1018 def build(self, input_shape=None):1019 if self.built:1020 return1021 self.built = True1022 if getattr(self, "data2vec_vision", None) is not None:1023 with tf.name_scope(self.data2vec_vision.name):1024 self.data2vec_vision.build(None)1025 1026 1027@add_start_docstrings(1028 """1029 Data2VecVision Model transformer with an image classification head on top (a linear layer on top of the average of1030 the final hidden states of the patch tokens) e.g. for ImageNet.1031 """,1032 DATA2VEC_VISION_START_DOCSTRING,1033)1034class TFData2VecVisionForImageClassification(TFData2VecVisionPreTrainedModel, TFSequenceClassificationLoss):1035 def __init__(self, config: Data2VecVisionConfig, *inputs, **kwargs):1036 super().__init__(config, *inputs, **kwargs)1037 1038 self.num_labels = config.num_labels1039 self.data2vec_vision = TFData2VecVisionMainLayer(config, add_pooling_layer=True, name="data2vec_vision")1040 1041 # Classifier head1042 self.classifier = keras.layers.Dense(1043 units=config.num_labels,1044 kernel_initializer=get_initializer(config.initializer_range),1045 name="classifier",1046 )1047 self.config = config1048 1049 @unpack_inputs1050 @add_start_docstrings_to_model_forward(DATA2VEC_VISION_INPUTS_DOCSTRING)1051 @add_code_sample_docstrings(1052 checkpoint=_IMAGE_CLASS_CHECKPOINT,1053 output_type=TFSequenceClassifierOutput,1054 config_class=_CONFIG_FOR_DOC,1055 expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,1056 )1057 def call(1058 self,1059 pixel_values: TFModelInputType | None = None,1060 head_mask: np.ndarray | tf.Tensor | None = None,1061 output_attentions: bool | None = None,1062 output_hidden_states: bool | None = None,1063 return_dict: bool | None = None,1064 labels: np.ndarray | tf.Tensor | None = None,1065 training: bool | None = False,1066 ) -> TFSequenceClassifierOutput | tuple:1067 r"""1068 labels (`tf.Tensor` or `np.ndarray` of shape `(batch_size,)`, *optional*):1069 Labels for computing the image classification/regression loss. Indices should be in `[0, ...,1070 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If1071 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).1072 """1073 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1074 1075 outputs = self.data2vec_vision(1076 pixel_values=pixel_values,1077 head_mask=head_mask,1078 output_attentions=output_attentions,1079 output_hidden_states=output_hidden_states,1080 return_dict=return_dict,1081 training=training,1082 )1083 1084 pooled_output = outputs.pooler_output if return_dict else outputs[1]1085 logits = self.classifier(pooled_output)1086 loss = None if labels is None else self.hf_compute_loss(labels=labels, logits=logits)1087 1088 if not return_dict:1089 output = (logits,) + outputs[2:]1090 return ((loss,) + output) if loss is not None else output1091 1092 return TFSequenceClassifierOutput(1093 loss=loss,1094 logits=logits,1095 hidden_states=outputs.hidden_states,1096 attentions=outputs.attentions,1097 )1098 1099 def build(self, input_shape=None):1100 if self.built:1101 return1102 self.built = True1103 if getattr(self, "data2vec_vision", None) is not None:1104 with tf.name_scope(self.data2vec_vision.name):1105 self.data2vec_vision.build(None)1106 if getattr(self, "classifier", None) is not None:1107 with tf.name_scope(self.classifier.name):1108 self.classifier.build([None, None, self.config.hidden_size])1109 1110 1111class TFData2VecVisionConvModule(keras.layers.Layer):1112 """1113 A convolutional block that bundles conv/norm/activation layers. This block simplifies the usage of convolution1114 layers, which are commonly used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU).1115 1116 Based on OpenMMLab's implementation, found in https://github.com/open-mmlab/mmsegmentation.1117 """1118 1119 def __init__(1120 self,1121 in_channels: int,1122 out_channels: int,1123 kernel_size: int | tuple[int, int],1124 padding: str = "valid",1125 bias: bool = False,1126 dilation: int | tuple[int, int] = 1,1127 **kwargs,1128 ) -> None:1129 super().__init__(**kwargs)1130 self.conv = keras.layers.Conv2D(1131 filters=out_channels,1132 kernel_size=kernel_size,1133 padding=padding,1134 use_bias=bias,1135 dilation_rate=dilation,1136 name="conv",1137 )1138 self.bn = keras.layers.BatchNormalization(name="bn", momentum=0.9, epsilon=1e-5)1139 self.activation = tf.nn.relu1140 self.in_channels = in_channels1141 self.out_channels = out_channels1142 1143 def call(self, input: tf.Tensor) -> tf.Tensor:1144 output = self.conv(input)1145 output = self.bn(output)1146 output = self.activation(output)1147 return output1148 1149 def build(self, input_shape=None):1150 if self.built:1151 return1152 self.built = True1153 if getattr(self, "conv", None) is not None:1154 with tf.name_scope(self.conv.name):1155 self.conv.build([None, None, None, self.in_channels])1156 if getattr(self, "bn", None) is not None:1157 with tf.name_scope(self.bn.name):1158 self.bn.build((None, None, None, self.out_channels))1159 1160 1161class TFAdaptiveAvgPool2D(keras.layers.Layer):1162 def __init__(self, output_dims: tuple[int, int], input_ordering: str = "NHWC", **kwargs):1163 super().__init__(**kwargs)1164 self.output_dims = output_dims1165 self.input_ordering = input_ordering1166 if input_ordering not in ("NCHW", "NHWC"):1167 raise ValueError("Unrecognized input_ordering, should be 'NCHW' or 'NHWC'!")1168 self.h_axis = input_ordering.index("H")1169 self.w_axis = input_ordering.index("W")1170 1171 def pseudo_1d_pool(self, inputs: tf.Tensor, h_pooling: bool):1172 # Figure out which axis we're pooling on1173 if h_pooling:1174 axis = self.h_axis1175 output_dim = self.output_dims[0]1176 else:1177 axis = self.w_axis1178 output_dim = self.output_dims[1]1179 input_dim = inputs.shape[axis]1180 1181 # Figure out the potential pooling windows1182 # This is the key idea - the torch op always uses only two1183 # consecutive pooling window sizes, like 3 and 4. Therefore,1184 # if we pool with both possible sizes, we simply need to gather1185 # the 'correct' pool at each position to reimplement the torch op.1186 small_window = math.ceil(input_dim / output_dim)1187 big_window = small_window + 11188 if h_pooling:1189 output_dim = self.output_dims[0]1190 small_window_shape = (small_window, 1)1191 big_window_shape = (big_window, 1)1192 else:1193 output_dim = self.output_dims[1]1194 small_window_shape = (1, small_window)1195 big_window_shape = (1, big_window)1196 1197 # For resizes to 1, or integer resizes, we can take quick shortcuts1198 if output_dim == input_dim:1199 return inputs1200 elif output_dim == 1: