tf-models-nightly 2.18.0.dev20240911__py2.py3-none-any.whl → 2.18.0.dev20240912__py2.py3-none-any.whl

official/vision/configs/backbones.py

@@ -52,6 +52,8 @@ class VisionTransformer(hyperparams.Config):
   layer_scale_init_value: float = 0.0
   # Transformer encoder spatial partition dimensions.
   transformer_partition_dims: Optional[Tuple[int, int, int, int]] = None
+  # If True, output attention scores.
+  output_attention_scores: bool = False
 
 
 @dataclasses.dataclass

official/vision/modeling/backbones/vit.py

@@ -108,22 +108,25 @@ class TokenLayer(layers.Layer):
 class Encoder(layers.Layer):
   """Transformer Encoder."""
 
-  def __init__(self,
-               num_layers,
-               mlp_dim,
-               num_heads,
-               dropout_rate=0.1,
-               attention_dropout_rate=0.1,
-               kernel_regularizer=None,
-               inputs_positions=None,
-               init_stochastic_depth_rate=0.0,
-               kernel_initializer='glorot_uniform',
-               add_pos_embed=True,
-               pos_embed_origin_shape=None,
-               pos_embed_target_shape=None,
-               layer_scale_init_value=0.0,
-               transformer_partition_dims=None,
-               **kwargs):
+  def __init__(
+      self,
+      num_layers,
+      mlp_dim,
+      num_heads,
+      dropout_rate=0.1,
+      attention_dropout_rate=0.1,
+      kernel_regularizer=None,
+      inputs_positions=None,
+      init_stochastic_depth_rate=0.0,
+      kernel_initializer='glorot_uniform',
+      add_pos_embed=True,
+      pos_embed_origin_shape=None,
+      pos_embed_target_shape=None,
+      layer_scale_init_value=0.0,
+      transformer_partition_dims=None,
+      output_attention_scores=False,
+      **kwargs,
+  ):
     super().__init__(**kwargs)
     self._num_layers = num_layers
     self._mlp_dim = mlp_dim
@@ -139,6 +142,7 @@ class Encoder(layers.Layer):
     self._pos_embed_target_shape = pos_embed_target_shape
     self._layer_scale_init_value = layer_scale_init_value
     self._transformer_partition_dims = transformer_partition_dims
+    self._output_attention_scores = output_attention_scores
 
   def build(self, input_shape):
     if self._add_pos_embed:
@@ -163,10 +167,13 @@ class Encoder(layers.Layer):
           kernel_initializer=self._kernel_initializer,
           norm_first=True,
           stochastic_depth_drop_rate=nn_layers.get_stochastic_depth_rate(
-              self._init_stochastic_depth_rate, i + 1, self._num_layers),
+              self._init_stochastic_depth_rate, i + 1, self._num_layers
+          ),
           norm_epsilon=1e-6,
           layer_scale_init_value=self._layer_scale_init_value,
-          transformer_partition_dims=self._transformer_partition_dims)
+          transformer_partition_dims=self._transformer_partition_dims,
+          return_attention_scores=self._output_attention_scores,
+      )
       self._encoder_layers.append(encoder_layer)
     self._norm = layers.LayerNormalization(epsilon=1e-6)
     super().build(input_shape)
@@ -177,9 +184,16 @@ class Encoder(layers.Layer):
       x = self._pos_embed(x, inputs_positions=self._inputs_positions)
     x = self._dropout(x, training=training)
 
+    attention_scores = None  # Needed to suppress undefined-variable warning.
     for encoder_layer in self._encoder_layers:
-      x = encoder_layer(x, training=training)
+      if self._output_attention_scores:
+        x, attention_scores = encoder_layer(x, training=training)
+      else:
+        x = encoder_layer(x, training=training)
     x = self._norm(x)
+
+    if self._output_attention_scores:
+      return x, attention_scores
     return x
 
   def get_config(self):
@@ -199,6 +213,7 @@ class Encoder(layers.Layer):
         'pos_embed_target_shape': self._pos_embed_target_shape,
         'layer_scale_init_value': self._layer_scale_init_value,
         'transformer_partition_dims': self._transformer_partition_dims,
+        'output_attention_scores': self._output_attention_scores,
     }
     config.update(updates)
     return config
@@ -227,6 +242,7 @@ class VisionTransformer(tf_keras.Model):
       pos_embed_shape: Optional[Tuple[int, int]] = None,
       layer_scale_init_value: float = 0.0,
       transformer_partition_dims: Optional[Tuple[int, int, int, int]] = None,
+      output_attention_scores: bool = False,
   ):
     """VisionTransformer initialization function."""
     self._mlp_dim = mlp_dim
@@ -265,20 +281,29 @@ class VisionTransformer(tf_keras.Model):
     if pooler == 'token':
       x = TokenLayer(name='cls')(x)
 
-    x = Encoder(
+    encoder_output = Encoder(
         num_layers=num_layers,
         mlp_dim=mlp_dim,
         num_heads=num_heads,
         dropout_rate=dropout_rate,
         attention_dropout_rate=attention_dropout_rate,
         kernel_regularizer=kernel_regularizer,
-        kernel_initializer='glorot_uniform' if original_init else dict(
-            class_name='TruncatedNormal', config=dict(stddev=.02)),
+        kernel_initializer='glorot_uniform'
+        if original_init
+        else dict(class_name='TruncatedNormal', config=dict(stddev=0.02)),
         init_stochastic_depth_rate=init_stochastic_depth_rate,
         pos_embed_origin_shape=pos_embed_shape,
         pos_embed_target_shape=pos_embed_target_shape,
-        layer_scale_init_value=layer_scale_init_value)(
-            x)
+        layer_scale_init_value=layer_scale_init_value,
+        output_attention_scores=output_attention_scores,
+    )(x)
+
+    endpoints = {}
+    if output_attention_scores:
+      x, attention_scores = encoder_output
+      endpoints['attention_scores'] = attention_scores
+    else:
+      x = encoder_output
 
     if pooler == 'token':
       output_feature = x[:, 1:]
@@ -292,7 +317,6 @@ class VisionTransformer(tf_keras.Model):
     else:
       raise ValueError(f'unrecognized pooler type: {pooler}')
 
-    endpoints = {}
     if output_2d_feature_maps:
       # Use the closest feature level.
       feat_level = round(math.log2(patch_size))
@@ -376,4 +400,6 @@ def build_vit(input_specs,
       output_2d_feature_maps=backbone_cfg.output_2d_feature_maps,
       layer_scale_init_value=backbone_cfg.layer_scale_init_value,
       pos_embed_shape=backbone_cfg.pos_embed_shape,
-      transformer_partition_dims=backbone_cfg.transformer_partition_dims)
+      transformer_partition_dims=backbone_cfg.transformer_partition_dims,
+      output_attention_scores=backbone_cfg.output_attention_scores,
+  )

official/vision/modeling/backbones/vit_test.py

@@ -95,6 +95,17 @@ class VisionTransformerTest(parameterized.TestCase, tf.test.TestCase):
     output = network(inputs)['pre_logits']
     self.assertEqual(output.shape, [1, 1, 1, 768])
 
+  def test_attention_scores(self):
+    tf_keras.backend.set_image_data_format('channels_last')
+    input_specs = tf_keras.layers.InputSpec(shape=[2, 224, 224, 3])
+    network = vit.VisionTransformer(
+        input_specs=input_specs, output_attention_scores=True
+    )
+
+    inputs = tf_keras.Input(shape=(224, 224, 3), batch_size=1)
+    outputs = network(inputs)
+    self.assertEqual(outputs['attention_scores'].shape, [1, 12, 197, 197])
+
 
 if __name__ == '__main__':
   tf.test.main()

{tf_models_nightly-2.18.0.dev20240911.dist-info → tf_models_nightly-2.18.0.dev20240912.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.18.0.dev20240911
+Version: 2.18.0.dev20240912
 Summary: TensorFlow Official Models
 Home-page: https://github.com/tensorflow/models
 Author: Google Inc.

{tf_models_nightly-2.18.0.dev20240911.dist-info → tf_models_nightly-2.18.0.dev20240912.dist-info}/RECORD

@@ -971,7 +971,7 @@ official/vision/registry_imports.py,sha256=__tuPecJUjyfXgiFst8ZJJT5OljeleDFT7c5i
 official/vision/train.py,sha256=8h7lbaC6WCgF7XLvr2tG3-hLhbtBaoP3a6bAt3E8QeU,4005
 official/vision/train_spatial_partitioning.py,sha256=xpEusyM-fEPVGyBuzlkRgsYGaPAobjypSjpapdj0-ec,5735
 official/vision/configs/__init__.py,sha256=-iKVbGCvFMGSHMC89utzGXvZ83BhW6JnbEfS38lbW3M,1045
-official/vision/configs/backbones.py,sha256=hlMwjUbpoAk_MPezXFcaAkDTBqyu7QGWTfrkeIAIG-E,5823
+official/vision/configs/backbones.py,sha256=BNvcKNqX1wmYT-OzZ_34rBLeFRK579kiWYh9PGV-aYw,5901
 official/vision/configs/backbones_3d.py,sha256=0lJsUzeYmuC5xiosOwrqlmgR1gkOa4tpSaxDbYYU7FE,3614
 official/vision/configs/common.py,sha256=arlsSF6_Q7Ng8WxqVpPqPHgAwCnGdskHs-DAQmqzu5I,6566
 official/vision/configs/decoders.py,sha256=_wG6MH1RzYuhMrvJu5menR7gDvklEXDSxlHJwzVF6H4,2080
@@ -1079,9 +1079,9 @@ official/vision/modeling/backbones/spinenet.py,sha256=FOCafyw_ZVIY76gzpiY8Al4mXr
 official/vision/modeling/backbones/spinenet_mobile.py,sha256=x2gcLs-caun6M9SYCgG9sKAm9YE3g6MfdT5YyvsBboA,20784
 official/vision/modeling/backbones/spinenet_mobile_test.py,sha256=Me6bQCfu1T3gliPsRj-MuMF4hREZb7SRPpPDJS_9LzY,3958
 official/vision/modeling/backbones/spinenet_test.py,sha256=Xf2N2IiN-x_cnL5_p22LPy06FjqG6XpCuVCySwcZUDE,4734
-official/vision/modeling/backbones/vit.py,sha256=-ErjrdAVH_lRZO320LjcP4PE4zy_vGJdlUqTYZI9Tvc,14438
+official/vision/modeling/backbones/vit.py,sha256=-qztjnzym-UO37nJCuSjlqCzkb9VPrXS2ZXs6CcbHto,15181
 official/vision/modeling/backbones/vit_specs.py,sha256=fJv6xLkpViz8W-ovBH6j76AWb09X2P2-6vDAkdH9Ezs,2412
-official/vision/modeling/backbones/vit_test.py,sha256=JzhzEaOfq1oQhTcfZF40S3iId1J8HsQvphoznAzo-MY,3463
+official/vision/modeling/backbones/vit_test.py,sha256=K4y-BPqCCcdqPZUSUuFazQwG_UPVJGn-XWI0FGtazhQ,3901
 official/vision/modeling/decoders/__init__.py,sha256=9u1sY2dRGvnCY678wBo4SjKOI8ywVYn_JvtwohhRlYM,815
 official/vision/modeling/decoders/aspp.py,sha256=tbHitR7pUvA5wiXsVpK904GjOOZLtXuJLq_FYlPLwTE,8610
 official/vision/modeling/decoders/aspp_test.py,sha256=w99TQPM2jaOo18XdUcZz_iWt6hpC_2WDAVq-cdbFWR8,3011
@@ -1222,9 +1222,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=nc6A9K53OGqF25xN5St8EiWvdVbda
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
 tensorflow_models/uplift/__init__.py,sha256=mqfa55gweOdpKoaQyid4A_4u7xw__FcQeSIF0k_pYmI,999
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.18.0.dev20240911.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.18.0.dev20240911.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.18.0.dev20240911.dist-info/METADATA,sha256=nqTkSETn7Mnd-KUOpzZerGxqtI6pMreziyNvlcP-QKg,1432
-tf_models_nightly-2.18.0.dev20240911.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.18.0.dev20240911.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.18.0.dev20240911.dist-info/RECORD,,
+tf_models_nightly-2.18.0.dev20240912.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.18.0.dev20240912.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.18.0.dev20240912.dist-info/METADATA,sha256=DSdv3ZNz6oi2xj_9C5HPYHoh0h1dxv8tz_LER44-4Ms,1432
+tf_models_nightly-2.18.0.dev20240912.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.18.0.dev20240912.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.18.0.dev20240912.dist-info/RECORD,,