tf-models-nightly 2.17.0.dev20240617__py2.py3-none-any.whl → 2.20.0.dev20251220__py2.py3-none-any.whl

official/vision/modeling/backbones/mobilenet.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,8 +17,6 @@
 import dataclasses
 from typing import Any
 
-# Import libraries
-
 from absl import logging
 import tensorflow as tf, tf_keras
 
@@ -766,6 +764,77 @@ def _mnv4_conv_medium_block_specs():
   }
 
 
+def _mnv4_conv_medium_seg_block_specs():
+  """Tailored MobileNetV4ConvMedium for dense prediction, e.g. segmentation."""
+
+  def convbn(kernel_size, strides, filters, output=False):
+    return BlockSpec(
+        block_fn='convbn',
+        kernel_size=kernel_size,
+        filters=filters,
+        strides=strides,
+        is_output=output,
+    )
+
+  def fused_ib(kernel_size, strides, filters, output=False):
+    return BlockSpec(
+        block_fn='fused_ib',
+        kernel_size=kernel_size,
+        filters=filters,
+        strides=strides,
+        expand_ratio=4.0,
+        is_output=output,
+    )
+
+  def uib(
+      start_dw_ks, middle_dw_ks, strides, filters, expand_ratio, output=False
+  ):
+    return BlockSpec(
+        block_fn='uib',
+        start_dw_kernel_size=start_dw_ks,
+        middle_dw_kernel_size=middle_dw_ks,
+        filters=filters,
+        strides=strides,
+        expand_ratio=expand_ratio,
+        use_layer_scale=False,
+        is_output=output,
+    )
+
+  blocks = [
+      convbn(3, 2, 32),
+      fused_ib(3, 2, 48, output=True),
+      # 3rd stage
+      uib(3, 5, 2, 80, 4.0),
+      uib(3, 3, 1, 80, 2.0, output=True),
+      # 4th stage
+      uib(3, 5, 2, 160, 6.0),
+      uib(3, 3, 1, 160, 4.0),
+      uib(3, 3, 1, 160, 4.0),
+      uib(3, 5, 1, 160, 4.0),
+      uib(3, 3, 1, 160, 4.0),
+      uib(3, 0, 1, 160, 4.0),
+      uib(3, 0, 1, 160, 4.0, output=True),
+      # 5th stage
+      uib(5, 5, 2, 256, 6.0),
+      uib(5, 5, 1, 128, 4.0),
+      uib(3, 5, 1, 128, 4.0),
+      uib(3, 5, 1, 128, 4.0),
+      uib(3, 0, 1, 128, 4.0),
+      uib(3, 5, 1, 128, 2.0),
+      uib(5, 5, 1, 128, 4.0),
+      uib(5, 0, 1, 128, 2.0, output=False),
+      # FC layers
+      convbn(1, 1, 448, output=True),
+      BlockSpec(block_fn='gpooling', is_output=False),
+      convbn(1, 1, 1280),
+  ]
+  return {
+      'spec_name': 'MobileNetV4ConvMediumSeg',
+      'block_spec_schema': block_spec_field_list(),
+      'block_specs': block_spec_values_to_list(blocks),
+  }
+
+
 MNV4ConvLarge_BLOCK_SPECS = {
     'spec_name': 'MobileNetV4ConvLarge',
     'block_spec_schema': [
@@ -1077,6 +1146,7 @@ SUPPORTED_SPECS_MAP = {
     'MobileNetV4ConvLarge': MNV4ConvLarge_BLOCK_SPECS,
     'MobileNetV4HybridMedium': _mnv4_hybrid_medium_block_specs(),
     'MobileNetV4HybridLarge': _mnv4_hybrid_large_block_specs(),
+    'MobileNetV4ConvMediumSeg': _mnv4_conv_medium_seg_block_specs(),
 }
 
 

official/vision/modeling/backbones/mobilenet_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,8 +17,6 @@
 import itertools
 import math
 
-# Import libraries
-
 from absl.testing import parameterized
 import tensorflow as tf, tf_keras
 
@@ -43,6 +41,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
       'MobileNetV4ConvLarge',
       'MobileNetV4HybridMedium',
       'MobileNetV4HybridLarge',
+      'MobileNetV4ConvMediumSeg',
   )
   def test_serialize_deserialize(self, model_id):
     # Create a network object that sets all of its config options.
@@ -96,6 +95,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV4ConvLarge',
               'MobileNetV4HybridMedium',
               'MobileNetV4HybridLarge',
+              'MobileNetV4ConvMediumSeg',
           ],
       )
   )
@@ -126,6 +126,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV4ConvLarge',
               'MobileNetV4HybridMedium',
               'MobileNetV4HybridLarge',
+              'MobileNetV4ConvMediumSeg',
           ],
           [32, 224],
       )
@@ -153,6 +154,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
         'MobileNetV4ConvLarge': [48, 96, 192, 512],
         'MobileNetV4HybridMedium': [48, 80, 160, 256],
         'MobileNetV4HybridLarge': [48, 96, 192, 512],
+        'MobileNetV4ConvMediumSeg': [48, 80, 160, 448],
     }
 
     network = mobilenet.MobileNet(model_id=model_id,
@@ -184,6 +186,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV4ConvLarge',
               'MobileNetV4HybridMedium',
               'MobileNetV4HybridLarge',
+              'MobileNetV4ConvMediumSeg',
           ],
           [32, 224],
       )
@@ -211,6 +214,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
         'MobileNetV4ConvLarge': [None, None, None, None],
         'MobileNetV4HybridMedium': [None, None, None, None],
         'MobileNetV4HybridLarge': [None, None, None, None],
+        'MobileNetV4ConvMediumSeg': [None, None, None, None],
     }
     network = mobilenet.MobileNet(model_id=model_id,
                                   filter_size_scale=1.0,
@@ -247,6 +251,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV4ConvLarge',
               'MobileNetV4HybridMedium',
               'MobileNetV4HybridLarge',
+              'MobileNetV4ConvMediumSeg',
           ],
           [1.0, 0.75],
       )
@@ -285,6 +290,8 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
         ('MobileNetV4HybridMedium', 0.75): 6072584,
         ('MobileNetV4HybridLarge', 1.0): 36648024,
         ('MobileNetV4HybridLarge', 0.75): 21598064,
+        ('MobileNetV4ConvMediumSeg', 1.0): 3787024,
+        ('MobileNetV4ConvMediumSeg', 0.75): 2302536,
     }
 
     input_size = 224
@@ -314,6 +321,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV4ConvLarge',
               'MobileNetV4HybridMedium',
               'MobileNetV4HybridLarge',
+              'MobileNetV4ConvMediumSeg',
           ],
           [8, 16, 32],
       )
@@ -340,6 +348,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
         'MobileNetV4ConvLarge': 512,
         'MobileNetV4HybridMedium': 256,
         'MobileNetV4HybridLarge': 512,
+        'MobileNetV4ConvMediumSeg': 448,
     }
 
     network = mobilenet.MobileNet(

official/vision/modeling/backbones/resnet.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,7 +16,6 @@
 
 from typing import Callable, Optional
 
-# Import libraries
 import tensorflow as tf, tf_keras
 
 from official.modeling import hyperparams

official/vision/modeling/backbones/resnet_3d.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,7 +15,6 @@
 """Contains definitions of 3D Residual Networks."""
 from typing import Callable, List, Tuple, Optional
 
-# Import libraries
 import tensorflow as tf, tf_keras
 
 from official.modeling import hyperparams

official/vision/modeling/backbones/resnet_3d_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,7 +14,6 @@
 
 """Tests for resnet."""
 
-# Import libraries
 from absl.testing import parameterized
 import tensorflow as tf, tf_keras
 

official/vision/modeling/backbones/resnet_deeplab.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,10 +14,11 @@
 
 """Contains definitions of Residual Networks with Deeplab modifications."""
 
-from typing import Callable, Optional, Tuple, List
+import math
+from typing import Callable, List, Optional, Tuple
 
-import numpy as np
 import tensorflow as tf, tf_keras
+
 from official.modeling import hyperparams
 from official.modeling import tf_utils
 from official.vision.modeling.backbones import factory
@@ -238,7 +239,7 @@ class DilatedResNet(tf_keras.Model):
     else:
       x = layers.MaxPool2D(pool_size=3, strides=2, padding='same')(x)
 
-    normal_resnet_stage = int(np.math.log2(self._output_stride)) - 2
+    normal_resnet_stage = int(math.log2(self._output_stride)) - 2
 
     endpoints = {}
     for i in range(normal_resnet_stage + 1):

official/vision/modeling/backbones/resnet_deeplab_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,7 +14,8 @@
 
 """Tests for resnet_deeplab models."""
 
-# Import libraries
+import math
+
 from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf, tf_keras
@@ -46,10 +47,15 @@ class ResNetTest(parameterized.TestCase, tf.test.TestCase):
     inputs = tf_keras.Input(shape=(input_size, input_size, 3), batch_size=1)
     endpoints = network(inputs)
     print(endpoints)
-    self.assertAllEqual([
-        1, input_size / output_stride, input_size / output_stride,
-        512 * endpoint_filter_scale
-    ], endpoints[str(int(np.math.log2(output_stride)))].shape.as_list())
+    self.assertAllEqual(
+        [
+            1,
+            input_size / output_stride,
+            input_size / output_stride,
+            512 * endpoint_filter_scale,
+        ],
+        endpoints[str(int(math.log2(output_stride)))].shape.as_list(),
+    )
 
   @parameterized.parameters(
       ('v0', None, 0.0, False, False),
@@ -82,10 +88,15 @@ class ResNetTest(parameterized.TestCase, tf.test.TestCase):
     inputs = tf_keras.Input(shape=(input_size, input_size, 3), batch_size=1)
     endpoints = network(inputs)
     print(endpoints)
-    self.assertAllEqual([
-        1, input_size / output_stride, input_size / output_stride,
-        512 * endpoint_filter_scale
-    ], endpoints[str(int(np.math.log2(output_stride)))].shape.as_list())
+    self.assertAllEqual(
+        [
+            1,
+            input_size / output_stride,
+            input_size / output_stride,
+            512 * endpoint_filter_scale,
+        ],
+        endpoints[str(int(math.log2(output_stride)))].shape.as_list(),
+    )
 
   @combinations.generate(
       combinations.combine(

official/vision/modeling/backbones/resnet_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,7 +14,6 @@
 
 """Tests for resnet."""
 
-# Import libraries
 from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf, tf_keras

official/vision/modeling/backbones/resnet_unet.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,7 +16,6 @@
 
 from typing import Callable, Optional, List
 
-# Import libraries
 import tensorflow as tf, tf_keras
 
 from official.modeling import hyperparams

official/vision/modeling/backbones/resnet_unet_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,8 +14,6 @@
 
 """Tests for resnet."""
 
-# Import libraries
-
 from absl.testing import parameterized
 import tensorflow as tf, tf_keras
 

official/vision/modeling/backbones/revnet.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,7 +15,6 @@
 """Contains definitions of RevNet."""
 
 from typing import Any, Callable, Dict, Optional
-# Import libraries
 import tensorflow as tf, tf_keras
 from official.modeling import hyperparams
 from official.modeling import tf_utils

official/vision/modeling/backbones/revnet_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,7 +14,6 @@
 
 """Tests for RevNet."""
 
-# Import libraries
 from absl.testing import parameterized
 import tensorflow as tf, tf_keras
 

official/vision/modeling/backbones/spinenet.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,8 +17,6 @@
 import math
 from typing import Any, List, Optional, Tuple
 
-# Import libraries
-
 from absl import logging
 import tensorflow as tf, tf_keras
 

official/vision/modeling/backbones/spinenet_mobile.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -30,8 +30,6 @@
 import math
 from typing import Any, List, Optional, Tuple
 
-# Import libraries
-
 from absl import logging
 import tensorflow as tf, tf_keras
 

official/vision/modeling/backbones/spinenet_mobile_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -27,7 +27,6 @@
 # limitations under the License.
 # ==============================================================================
 """Tests for SpineNet."""
-# Import libraries
 from absl.testing import parameterized
 import tensorflow as tf, tf_keras
 

official/vision/modeling/backbones/spinenet_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,7 +13,6 @@
 # limitations under the License.
 
 """Tests for SpineNet."""
-# Import libraries
 from absl.testing import parameterized
 import tensorflow as tf, tf_keras
 

official/vision/modeling/backbones/vit.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -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_specs.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

official/vision/modeling/backbones/vit_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -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()

official/vision/modeling/classification_model.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,7 +15,6 @@
 """Build classification models."""
 
 from typing import Any, Mapping, Optional
-# Import libraries
 import tensorflow as tf, tf_keras
 
 layers = tf_keras.layers

official/vision/modeling/classification_model_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,7 +14,6 @@
 
 """Tests for classification network."""
 
-# Import libraries
 from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf, tf_keras