Internal change

PiperOrigin-RevId: 331232905

official/vision/beta/modeling/backbones/efficientnet.py

@@ -189,7 +189,7 @@ class EfficientNet(tf.keras.Model):
       x = self._block_group(
           inputs=x, specs=specs, name='block_group_{}'.format(i))
       if specs.is_output:
-        endpoints[endpoint_level] = x
+        endpoints[str(endpoint_level)] = x
         endpoint_level += 1
 
     # Build output specs for downstream tasks.
@@ -209,7 +209,7 @@ class EfficientNet(tf.keras.Model):
     x = self._norm(
         axis=bn_axis, momentum=norm_momentum, epsilon=norm_epsilon)(
             x)
-    endpoints[endpoint_level] = tf_utils.get_activation(activation)(x)
+    endpoints[str(endpoint_level)] = tf_utils.get_activation(activation)(x)
 
     super(EfficientNet, self).__init__(
         inputs=inputs, outputs=endpoints, **kwargs)

official/vision/beta/modeling/backbones/efficientnet_test.py

@@ -35,13 +35,13 @@ class EfficientNetTest(parameterized.TestCase, tf.test.TestCase):
     endpoints = network(inputs)
 
     self.assertAllEqual([1, input_size / 2**2, input_size / 2**2, 24],
-                        endpoints[2].shape.as_list())
+                        endpoints['2'].shape.as_list())
     self.assertAllEqual([1, input_size / 2**3, input_size / 2**3, 40],
-                        endpoints[3].shape.as_list())
+                        endpoints['3'].shape.as_list())
     self.assertAllEqual([1, input_size / 2**4, input_size / 2**4, 112],
-                        endpoints[4].shape.as_list())
+                        endpoints['4'].shape.as_list())
     self.assertAllEqual([1, input_size / 2**5, input_size / 2**5, 320],
-                        endpoints[5].shape.as_list())
+                        endpoints['5'].shape.as_list())
 
   @parameterized.parameters('b0', 'b3', 'b6')
   def test_network_scaling(self, model_id):

official/vision/beta/modeling/backbones/resnet.py

@@ -152,7 +152,7 @@ class ResNet(tf.keras.Model):
           block_fn=block_fn,
           block_repeats=spec[2],
           name='block_group_l{}'.format(i + 2))
-      endpoints[i + 2] = x
+      endpoints[str(i + 2)] = x
 
     self._output_specs = {l: endpoints[l].get_shape() for l in endpoints}
 

official/vision/beta/modeling/backbones/resnet_test.py

@@ -54,16 +54,16 @@ class ResNetTest(parameterized.TestCase, tf.test.TestCase):
 
     self.assertAllEqual(
         [1, input_size / 2**2, input_size / 2**2, 64 * endpoint_filter_scale],
-        endpoints[2].shape.as_list())
+        endpoints['2'].shape.as_list())
     self.assertAllEqual(
         [1, input_size / 2**3, input_size / 2**3, 128 * endpoint_filter_scale],
-        endpoints[3].shape.as_list())
+        endpoints['3'].shape.as_list())
     self.assertAllEqual(
         [1, input_size / 2**4, input_size / 2**4, 256 * endpoint_filter_scale],
-        endpoints[4].shape.as_list())
+        endpoints['4'].shape.as_list())
     self.assertAllEqual(
         [1, input_size / 2**5, input_size / 2**5, 512 * endpoint_filter_scale],
-        endpoints[5].shape.as_list())
+        endpoints['5'].shape.as_list())
 
   @combinations.generate(
       combinations.combine(

official/vision/beta/modeling/backbones/revnet.py

@@ -130,7 +130,7 @@ class RevNet(tf.keras.Model):
           block_repeats=spec[2],
           batch_norm_first=(i != 0),  # Only skip on first block
           name='revblock_group_{}'.format(i + 2))
-      endpoints[i + 2] = x
+      endpoints[str(i + 2)] = x
 
     self._output_specs = {l: endpoints[l].get_shape() for l in endpoints}
 

official/vision/beta/modeling/backbones/revnet_test.py

@@ -40,16 +40,16 @@ class RevNetTest(parameterized.TestCase, tf.test.TestCase):
 
     self.assertAllEqual(
         [1, input_size / 2**2, input_size / 2**2, 128 * endpoint_filter_scale],
-        endpoints[2].shape.as_list())
+        endpoints['2'].shape.as_list())
     self.assertAllEqual(
         [1, input_size / 2**3, input_size / 2**3, 256 * endpoint_filter_scale],
-        endpoints[3].shape.as_list())
+        endpoints['3'].shape.as_list())
     self.assertAllEqual(
         [1, input_size / 2**4, input_size / 2**4, 512 * endpoint_filter_scale],
-        endpoints[4].shape.as_list())
+        endpoints['4'].shape.as_list())
     self.assertAllEqual(
         [1, input_size / 2**5, input_size / 2**5, 832 * endpoint_filter_scale],
-        endpoints[5].shape.as_list())
+        endpoints['5'].shape.as_list())
 
   @parameterized.parameters(1, 3, 4)
   def test_input_specs(self, input_dim):

official/vision/beta/modeling/backbones/spinenet.py

@@ -349,7 +349,7 @@ class SpineNet(tf.keras.Model):
             block_spec.level > self._max_level):
           raise ValueError('Output level is out of range [{}, {}]'.format(
               self._min_level, self._max_level))
-        endpoints[block_spec.level] = x
+        endpoints[str(block_spec.level)] = x
 
     return endpoints
 
@@ -365,14 +365,14 @@ class SpineNet(tf.keras.Model):
           kernel_initializer=self._kernel_initializer,
           kernel_regularizer=self._kernel_regularizer,
           bias_regularizer=self._bias_regularizer)(
-              net[level])
+              net[str(level)])
       x = self._norm(
           axis=self._bn_axis,
           momentum=self._norm_momentum,
           epsilon=self._norm_epsilon)(
               x)
       x = tf_utils.get_activation(self._activation_fn)(x)
-      endpoints[level] = x
+      endpoints[str(level)] = x
     return endpoints
 
   def _resample_with_alpha(self,

official/vision/beta/modeling/backbones/spinenet_test.py

@@ -54,10 +54,10 @@ class SpineNetTest(parameterized.TestCase, tf.test.TestCase):
     endpoints = model(inputs)
 
     for l in range(min_level, max_level + 1):
-      self.assertIn(l, endpoints.keys())
+      self.assertIn(str(l), endpoints.keys())
       self.assertAllEqual(
           [1, input_size / 2**l, input_size / 2**l, endpoints_num_filters],
-          endpoints[l].shape.as_list())
+          endpoints[str(l)].shape.as_list())
 
   def test_serialize_deserialize(self):
     # Create a network object that sets all of its config options.

official/vision/beta/modeling/decoders/fpn.py

@@ -98,30 +98,30 @@ class FPN(tf.keras.Model):
 
     # Get input feature pyramid from backbone.
     inputs = self._build_input_pyramid(input_specs, min_level)
-    backbone_max_level = min(max(inputs.keys()), max_level)
+    backbone_max_level = min(int(max(inputs.keys())), max_level)
 
     # Build lateral connections.
     feats_lateral = {}
     for level in range(min_level, backbone_max_level + 1):
-      feats_lateral[level] = conv2d(
+      feats_lateral[str(level)] = conv2d(
           filters=num_filters,
           kernel_size=1,
           padding='same',
           kernel_initializer=kernel_initializer,
           kernel_regularizer=kernel_regularizer,
           bias_regularizer=bias_regularizer)(
-              inputs[level])
+              inputs[str(level)])
 
     # Build top-down path.
-    feats = {backbone_max_level: feats_lateral[backbone_max_level]}
+    feats = {str(backbone_max_level): feats_lateral[str(backbone_max_level)]}
     for level in range(backbone_max_level - 1, min_level - 1, -1):
-      feats[level] = spatial_transform_ops.nearest_upsampling(
-          feats[level + 1], 2) + feats_lateral[level]
+      feats[str(level)] = spatial_transform_ops.nearest_upsampling(
+          feats[str(level + 1)], 2) + feats_lateral[str(level)]
 
     # TODO(xianzhi): consider to remove bias in conv2d.
     # Build post-hoc 3x3 convolution kernel.
     for level in range(min_level, backbone_max_level + 1):
-      feats[level] = conv2d(
+      feats[str(level)] = conv2d(
           filters=num_filters,
           strides=1,
           kernel_size=3,
@@ -129,15 +129,15 @@ class FPN(tf.keras.Model):
           kernel_initializer=kernel_initializer,
           kernel_regularizer=kernel_regularizer,
           bias_regularizer=bias_regularizer)(
-              feats[level])
+              feats[str(level)])
 
     # TODO(xianzhi): consider to remove bias in conv2d.
     # Build coarser FPN levels introduced for RetinaNet.
     for level in range(backbone_max_level + 1, max_level + 1):
-      feats_in = feats[level - 1]
+      feats_in = feats[str(level - 1)]
       if level > backbone_max_level + 1:
         feats_in = activation_fn(feats_in)
-      feats[level] = conv2d(
+      feats[str(level)] = conv2d(
           filters=num_filters,
           strides=2,
           kernel_size=3,
@@ -149,12 +149,12 @@ class FPN(tf.keras.Model):
 
     # Apply batch norm layers.
     for level in range(min_level, max_level + 1):
-      feats[level] = norm(
+      feats[str(level)] = norm(
           axis=bn_axis, momentum=norm_momentum, epsilon=norm_epsilon)(
-              feats[level])
+              feats[str(level)])
 
     self._output_specs = {
-        level: feats[level].get_shape()
+        str(level): feats[str(level)].get_shape()
         for level in range(min_level, max_level + 1)
     }
 
@@ -162,7 +162,7 @@ class FPN(tf.keras.Model):
 
   def _build_input_pyramid(self, input_specs, min_level):
     assert isinstance(input_specs, dict)
-    if min(input_specs.keys()) > min_level:
+    if min(input_specs.keys()) > str(min_level):
       raise ValueError(
           'Backbone min level should be less or equal to FPN min level')
 

official/vision/beta/modeling/decoders/fpn_test.py

@@ -47,10 +47,10 @@ class FPNTest(parameterized.TestCase, tf.test.TestCase):
     feats = network(endpoints)
 
     for level in range(min_level, max_level + 1):
-      self.assertIn(level, feats)
+      self.assertIn(str(level), feats)
       self.assertAllEqual(
           [1, input_size // 2**level, input_size // 2**level, 256],
-          feats[level].shape.as_list())
+          feats[str(level)].shape.as_list())
 
   def test_serialize_deserialize(self):
     # Create a network object that sets all of its config options.

official/vision/beta/modeling/heads/dense_prediction_heads.py

@@ -210,18 +210,18 @@ class RetinaNetHead(tf.keras.layers.Layer):
 
     Args:
       features: a dict of tensors
-        - key: `int`, the level of the multilevel features.
+        - key: `str`, the level of the multilevel features.
         - values: `Tensor`, the feature map tensors, whose shape is
             [batch, height_l, width_l, channels].
 
     Returns:
       scores: a dict of tensors which includes scores of the predictions.
-        - key: `int`, the level of the multilevel predictions.
+        - key: `str`, the level of the multilevel predictions.
         - values: `Tensor`, the box scores predicted from a particular feature
             level, whose shape is
             [batch, height_l, width_l, num_classes * num_anchors_per_location].
       boxes: a dict of tensors which includes coordinates of the predictions.
-        - key: `int`, the level of the multilevel predictions.
+        - key: `str`, the level of the multilevel predictions.
         - values: `Tensor`, the box scores predicted from a particular feature
             level, whose shape is
             [batch, height_l, width_l, 4 * num_anchors_per_location].
@@ -231,7 +231,7 @@ class RetinaNetHead(tf.keras.layers.Layer):
     for i, level in enumerate(
         range(self._config_dict['min_level'],
               self._config_dict['max_level'] + 1)):
-      this_level_features = features[level]
+      this_level_features = features[str(level)]
 
       # class net.
       x = this_level_features
@@ -239,7 +239,7 @@ class RetinaNetHead(tf.keras.layers.Layer):
         x = conv(x)
         x = norm(x)
         x = self._activation(x)
-      scores[level] = self._classifier(x)
+      scores[str(level)] = self._classifier(x)
 
       # box net.
       x = this_level_features
@@ -247,7 +247,7 @@ class RetinaNetHead(tf.keras.layers.Layer):
         x = conv(x)
         x = norm(x)
         x = self._activation(x)
-      boxes[level] = self._box_regressor(x)
+      boxes[str(level)] = self._box_regressor(x)
     return scores, boxes
 
   def get_config(self):
@@ -433,13 +433,13 @@ class RPNHead(tf.keras.layers.Layer):
     for i, level in enumerate(
         range(self._config_dict['min_level'],
               self._config_dict['max_level'] + 1)):
-      x = features[level]
+      x = features[str(level)]
       for conv, norm in zip(self._convs, self._norms[i]):
         x = conv(x)
         x = norm(x)
         x = self._activation(x)
-      scores[level] = self._classifier(x)
-      boxes[level] = self._box_regressor(x)
+      scores[str(level)] = self._classifier(x)
+      boxes[str(level)] = self._box_regressor(x)
     return scores, boxes
 
   def get_config(self):

official/vision/beta/modeling/heads/dense_prediction_heads_test.py

@@ -48,14 +48,14 @@ class RetinaNetHeadTest(parameterized.TestCase, tf.test.TestCase):
         bias_regularizer=None,
     )
     features = {
-        3: np.random.rand(2, 128, 128, 16),
-        4: np.random.rand(2, 64, 64, 16),
+        '3': np.random.rand(2, 128, 128, 16),
+        '4': np.random.rand(2, 64, 64, 16),
     }
     scores, boxes = retinanet_head(features)
-    self.assertAllEqual(scores[3].numpy().shape, [2, 128, 128, 9])
-    self.assertAllEqual(scores[4].numpy().shape, [2, 64, 64, 9])
-    self.assertAllEqual(boxes[3].numpy().shape, [2, 128, 128, 12])
-    self.assertAllEqual(boxes[4].numpy().shape, [2, 64, 64, 12])
+    self.assertAllEqual(scores['3'].numpy().shape, [2, 128, 128, 9])
+    self.assertAllEqual(scores['4'].numpy().shape, [2, 64, 64, 9])
+    self.assertAllEqual(boxes['3'].numpy().shape, [2, 128, 128, 12])
+    self.assertAllEqual(boxes['4'].numpy().shape, [2, 64, 64, 12])
 
   def test_serialize_deserialize(self):
     retinanet_head = dense_prediction_heads.RetinaNetHead(
@@ -104,14 +104,14 @@ class RpnHeadTest(parameterized.TestCase, tf.test.TestCase):
         bias_regularizer=None,
     )
     features = {
-        3: np.random.rand(2, 128, 128, 16),
-        4: np.random.rand(2, 64, 64, 16),
+        '3': np.random.rand(2, 128, 128, 16),
+        '4': np.random.rand(2, 64, 64, 16),
     }
     scores, boxes = rpn_head(features)
-    self.assertAllEqual(scores[3].numpy().shape, [2, 128, 128, 3])
-    self.assertAllEqual(scores[4].numpy().shape, [2, 64, 64, 3])
-    self.assertAllEqual(boxes[3].numpy().shape, [2, 128, 128, 12])
-    self.assertAllEqual(boxes[4].numpy().shape, [2, 64, 64, 12])
+    self.assertAllEqual(scores['3'].numpy().shape, [2, 128, 128, 3])
+    self.assertAllEqual(scores['4'].numpy().shape, [2, 64, 64, 3])
+    self.assertAllEqual(boxes['3'].numpy().shape, [2, 128, 128, 12])
+    self.assertAllEqual(boxes['4'].numpy().shape, [2, 64, 64, 12])
 
   def test_serialize_deserialize(self):
     rpn_head = dense_prediction_heads.RPNHead(

official/vision/beta/modeling/layers/detection_generator.py

@@ -563,24 +563,25 @@ class MultilevelDetectionGenerator(tf.keras.layers.Layer):
     boxes = []
     scores = []
     levels = list(raw_boxes.keys())
-    min_level = min(levels)
-    max_level = max(levels)
+    min_level = int(min(levels))
+    max_level = int(max(levels))
     for i in range(min_level, max_level + 1):
-      raw_boxes_i_shape = tf.shape(raw_boxes[i])
+      raw_boxes_i_shape = tf.shape(raw_boxes[str(i)])
       batch_size = raw_boxes_i_shape[0]
       num_anchors_per_locations = raw_boxes_i_shape[-1] // 4
-      num_classes = tf.shape(raw_scores[i])[-1] // num_anchors_per_locations
+      num_classes = tf.shape(
+          raw_scores[str(i)])[-1] // num_anchors_per_locations
 
       # Applies score transformation and remove the implicit background class.
       scores_i = tf.sigmoid(
-          tf.reshape(raw_scores[i], [batch_size, -1, num_classes]))
+          tf.reshape(raw_scores[str(i)], [batch_size, -1, num_classes]))
       scores_i = tf.slice(scores_i, [0, 0, 1], [-1, -1, -1])
 
       # Box decoding.
       # The anchor boxes are shared for all data in a batch.
       # One stage detector only supports class agnostic box regression.
-      anchor_boxes_i = tf.reshape(anchor_boxes[i], [batch_size, -1, 4])
-      raw_boxes_i = tf.reshape(raw_boxes[i], [batch_size, -1, 4])
+      anchor_boxes_i = tf.reshape(anchor_boxes[str(i)], [batch_size, -1, 4])
+      raw_boxes_i = tf.reshape(raw_boxes[str(i)], [batch_size, -1, 4])
       boxes_i = box_ops.decode_boxes(raw_boxes_i, anchor_boxes_i)
 
       # Box clipping.

official/vision/beta/modeling/layers/detection_generator_test.py

@@ -148,22 +148,25 @@ class MultilevelDetectionGeneratorTest(
         np.random.rand(84, num_classes) - 0.5) * 3  # random 84x3 outputs.
     box_outputs_all = np.random.rand(84, 4)  # random 84 boxes.
     class_outputs = {
-        4: tf.reshape(tf.convert_to_tensor(
-            cls_outputs_all[0:64], dtype=tf.float32),
-                      [1, 8, 8, num_classes]),
-        5: tf.reshape(tf.convert_to_tensor(
-            cls_outputs_all[64:80], dtype=tf.float32),
-                      [1, 4, 4, num_classes]),
-        6: tf.reshape(tf.convert_to_tensor(
-            cls_outputs_all[80:84], dtype=tf.float32),
-                      [1, 2, 2, num_classes]),
+        '4':
+            tf.reshape(
+                tf.convert_to_tensor(cls_outputs_all[0:64], dtype=tf.float32),
+                [1, 8, 8, num_classes]),
+        '5':
+            tf.reshape(
+                tf.convert_to_tensor(cls_outputs_all[64:80], dtype=tf.float32),
+                [1, 4, 4, num_classes]),
+        '6':
+            tf.reshape(
+                tf.convert_to_tensor(cls_outputs_all[80:84], dtype=tf.float32),
+                [1, 2, 2, num_classes]),
     }
     box_outputs = {
-        4: tf.reshape(tf.convert_to_tensor(
+        '4': tf.reshape(tf.convert_to_tensor(
             box_outputs_all[0:64], dtype=tf.float32), [1, 8, 8, 4]),
-        5: tf.reshape(tf.convert_to_tensor(
+        '5': tf.reshape(tf.convert_to_tensor(
             box_outputs_all[64:80], dtype=tf.float32), [1, 4, 4, 4]),
-        6: tf.reshape(tf.convert_to_tensor(
+        '6': tf.reshape(tf.convert_to_tensor(
             box_outputs_all[80:84], dtype=tf.float32), [1, 2, 2, 4]),
     }
     image_info = tf.constant([[[1000, 1000], [100, 100], [0.1, 0.1], [0, 0]]],

official/vision/beta/modeling/layers/roi_generator.py

@@ -95,7 +95,7 @@ def _multilevel_propose_rois(raw_boxes,
     roi_scores = []
     image_shape = tf.expand_dims(image_shape, axis=1)
     for level in sorted(raw_scores.keys()):
-      with tf.name_scope('level_%d' % level):
+      with tf.name_scope('level_%s' % level):
         _, feature_h, feature_w, num_anchors_per_location = (
             raw_scores[level].get_shape().as_list())
 

official/vision/beta/modeling/layers/roi_generator_test.py

@@ -30,12 +30,12 @@ class MultilevelProposeRoisTest(tf.test.TestCase):
          [[[2, 2, 4, 4], [3, 3, 6, 6]],
           [[3.1, 3.1, 6.1, 6.1], [1, 1, 8, 8]]]])
     rpn_boxes = {
-        2: tf.constant(rpn_boxes_np, dtype=tf.float32)
+        '2': tf.constant(rpn_boxes_np, dtype=tf.float32)
     }
     rpn_scores_np = np.array(
         [[[[0.6], [0.9]], [[0.2], [0.3]]], [[[0.1], [0.8]], [[0.3], [0.5]]]])
     rpn_scores = {
-        2: tf.constant(rpn_scores_np, dtype=tf.float32)
+        '2': tf.constant(rpn_scores_np, dtype=tf.float32)
     }
     anchor_boxes_np = np.array(
         [[[[0, 0, 10, 10], [0.01, 0.01, 9.9, 9.9]],
@@ -43,7 +43,7 @@ class MultilevelProposeRoisTest(tf.test.TestCase):
          [[[2, 2, 4, 4], [3, 3, 6, 6]],
           [[3.1, 3.1, 6.1, 6.1], [1, 1, 8, 8]]]])
     anchor_boxes = {
-        2: tf.constant(anchor_boxes_np, dtype=tf.float32)
+        '2': tf.constant(anchor_boxes_np, dtype=tf.float32)
     }
     image_shape = tf.constant([[20, 20], [20, 20]], dtype=tf.int32)
 
@@ -108,15 +108,15 @@ class MultilevelProposeRoisTest(tf.test.TestCase):
     rpn_boxes_2_np = np.array(
         [[[[0, 0, 10.01, 10.01]]], [[[2, 2, 4.5, 4.5]]]])
     rpn_boxes = {
-        2: tf.constant(rpn_boxes_1_np, dtype=tf.float32),
-        3: tf.constant(rpn_boxes_2_np, dtype=tf.float32),
+        '2': tf.constant(rpn_boxes_1_np, dtype=tf.float32),
+        '3': tf.constant(rpn_boxes_2_np, dtype=tf.float32),
     }
     rpn_scores_1_np = np.array(
         [[[[0.6], [0.9]], [[0.2], [0.3]]], [[[0.1], [0.8]], [[0.3], [0.5]]]])
     rpn_scores_2_np = np.array([[[[0.95]]], [[[0.99]]]])
     rpn_scores = {
-        2: tf.constant(rpn_scores_1_np, dtype=tf.float32),
-        3: tf.constant(rpn_scores_2_np, dtype=tf.float32),
+        '2': tf.constant(rpn_scores_1_np, dtype=tf.float32),
+        '3': tf.constant(rpn_scores_2_np, dtype=tf.float32),
     }
     anchor_boxes_1_np = np.array(
         [[[[0, 0, 10, 10], [0.01, 0.01, 9.99, 9.99]],
@@ -126,8 +126,8 @@ class MultilevelProposeRoisTest(tf.test.TestCase):
     anchor_boxes_2_np = np.array(
         [[[[0, 0, 10.01, 10.01]]], [[[2, 2, 4.5, 4.5]]]])
     anchor_boxes = {
-        2: tf.constant(anchor_boxes_1_np, dtype=tf.float32),
-        3: tf.constant(anchor_boxes_2_np, dtype=tf.float32),
+        '2': tf.constant(anchor_boxes_1_np, dtype=tf.float32),
+        '3': tf.constant(anchor_boxes_2_np, dtype=tf.float32),
     }
     image_shape = tf.constant([[20, 20], [20, 20]], dtype=tf.int32)
 
@@ -208,3 +208,6 @@ class MultilevelROIGeneratorTest(tf.test.TestCase):
         generator.get_config())
 
     self.assertAllEqual(generator.get_config(), new_generator.get_config())
+
+if __name__ == '__main__':
+  tf.test.main()

official/vision/beta/modeling/retinanet_model.py

@@ -65,19 +65,19 @@ class RetinaNetModel(tf.keras.Model):
         in the batch may be resized into different shapes before padding to the
         fixed size.
       anchor_boxes: a dict of tensors which includes multilevel anchors.
-        - key: `int`, the level of the multilevel predictions.
+        - key: `str`, the level of the multilevel predictions.
         - values: `Tensor`, the anchor coordinates of a particular feature
             level, whose shape is [height_l, width_l, num_anchors_per_location].
       training: `bool`, indicating whether it is in training mode.
 
     Returns:
       scores: a dict of tensors which includes scores of the predictions.
-        - key: `int`, the level of the multilevel predictions.
+        - key: `str`, the level of the multilevel predictions.
         - values: `Tensor`, the box scores predicted from a particular feature
             level, whose shape is
             [batch, height_l, width_l, num_classes * num_anchors_per_location].
       boxes: a dict of tensors which includes coordinates of the predictions.
-        - key: `int`, the level of the multilevel predictions.
+        - key: `str`, the level of the multilevel predictions.
         - values: `Tensor`, the box coordinates predicted from a particular
             feature level, whose shape is
             [batch, height_l, width_l, 4 * num_anchors_per_location].

official/vision/beta/modeling/retinanet_model_test.py

@@ -153,20 +153,20 @@ class RetinaNetTest(parameterized.TestCase, tf.test.TestCase):
       cls_outputs = model_outputs['cls_outputs']
       box_outputs = model_outputs['box_outputs']
       for level in range(min_level, max_level + 1):
-        self.assertIn(level, cls_outputs)
-        self.assertIn(level, box_outputs)
+        self.assertIn(str(level), cls_outputs)
+        self.assertIn(str(level), box_outputs)
         self.assertAllEqual([
             2,
             image_size[0] // 2**level,
             image_size[1] // 2**level,
             num_classes * num_anchors_per_location
-        ], cls_outputs[level].numpy().shape)
+        ], cls_outputs[str(level)].numpy().shape)
         self.assertAllEqual([
             2,
             image_size[0] // 2**level,
             image_size[1] // 2**level,
             4 * num_anchors_per_location
-        ], box_outputs[level].numpy().shape)
+        ], box_outputs[str(level)].numpy().shape)
     else:
       self.assertIn('detection_boxes', model_outputs)
       self.assertIn('detection_scores', model_outputs)
@@ -220,3 +220,4 @@ class RetinaNetTest(parameterized.TestCase, tf.test.TestCase):
 
 if __name__ == '__main__':
   tf.test.main()
+

official/vision/beta/ops/anchor.py

@@ -104,7 +104,7 @@ class Anchor(object):
       feat_size_y = tf.cast(self.image_size[0] / 2 ** level, tf.int32)
       feat_size_x = tf.cast(self.image_size[1] / 2 ** level, tf.int32)
       steps = feat_size_y * feat_size_x * self.anchors_per_location
-      unpacked_labels[level] = tf.reshape(
+      unpacked_labels[str(level)] = tf.reshape(
           labels[count:count + steps], [feat_size_y, feat_size_x, -1])
       count += steps
     return unpacked_labels
@@ -315,8 +315,8 @@ def build_anchor_generator(min_level, max_level, num_scales, aspect_ratios,
     scales.append(2**(scale / float(num_scales)))
   for level in range(min_level, max_level + 1):
     stride = 2**level
-    strides[level] = stride
-    anchor_sizes[level] = anchor_size * stride
+    strides[str(level)] = stride
+    anchor_sizes[str(level)] = anchor_size * stride
   anchor_gen = keras_cv.ops.AnchorGenerator(
       anchor_sizes=anchor_sizes,
       scales=scales,

official/vision/beta/ops/anchor_test.py

@@ -124,7 +124,7 @@ class AnchorTest(parameterized.TestCase, tf.test.TestCase):
 
     # Uses the first anchors as ground truth. The ground truth should map to
     # two anchors with two intermediate scales at the same location.
-    gt_boxes = anchor_boxes[3][0:1, 0, 0:4]
+    gt_boxes = anchor_boxes['3'][0:1, 0, 0:4]
     gt_classes = tf.constant([[ground_truth_class_id]], dtype=tf.float32)
     (cls_targets, box_targets, _,
      box_weights) = anchor_labeler.label_anchors(
@@ -137,7 +137,7 @@ class AnchorTest(parameterized.TestCase, tf.test.TestCase):
     box_weights = box_weights.numpy()
 
     anchor_locations = np.vstack(
-        np.where(cls_targets[min_level] > -1)).transpose()
+        np.where(cls_targets[str(min_level)] > -1)).transpose()
     self.assertAllClose(expected_anchor_locations, anchor_locations)
     # Two anchor boxes on min_level got matched to the gt_boxes.
     self.assertAllClose(tf.reduce_sum(box_weights), 2)