Merge branch 'master' of https://github.com/tensorflow/models into RTESuperGLUE

official/vision/beta/serving/semantic_segmentation_test.py

@@ -38,33 +38,9 @@ class SemanticSegmentationExportTest(tf.test.TestCase, parameterized.TestCase):
     return segmentation_module
 
   def _export_from_module(self, module, input_type, save_directory):
-    if input_type == 'image_tensor':
-      input_signature = tf.TensorSpec(shape=[None, 112, 112, 3], dtype=tf.uint8)
-      signatures = {
-          'serving_default':
-              module.inference_from_image_tensors.get_concrete_function(
-                  input_signature)
-      }
-    elif input_type == 'image_bytes':
-      input_signature = tf.TensorSpec(shape=[None], dtype=tf.string)
-      signatures = {
-          'serving_default':
-              module.inference_from_image_bytes.get_concrete_function(
-                  input_signature)
-      }
-    elif input_type == 'tf_example':
-      input_signature = tf.TensorSpec(shape=[None], dtype=tf.string)
-      signatures = {
-          'serving_default':
-              module.inference_from_tf_example.get_concrete_function(
-                  input_signature)
-      }
-    else:
-      raise ValueError('Unrecognized `input_type`')
-
-    tf.saved_model.save(module,
-                        save_directory,
-                        signatures=signatures)
+    signatures = module.get_inference_signatures(
+        {input_type: 'serving_default'})
+    tf.saved_model.save(module, save_directory, signatures=signatures)
 
   def _get_dummy_input(self, input_type):
     """Get dummy input for the given input type."""
@@ -95,17 +71,17 @@ class SemanticSegmentationExportTest(tf.test.TestCase, parameterized.TestCase):
   )
   def test_export(self, input_type='image_tensor'):
     tmp_dir = self.get_temp_dir()
-
     module = self._get_segmentation_module()
-    model = module.build_model()
 
     self._export_from_module(module, input_type, tmp_dir)
 
     self.assertTrue(os.path.exists(os.path.join(tmp_dir, 'saved_model.pb')))
-    self.assertTrue(os.path.exists(
-        os.path.join(tmp_dir, 'variables', 'variables.index')))
-    self.assertTrue(os.path.exists(
-        os.path.join(tmp_dir, 'variables', 'variables.data-00000-of-00001')))
+    self.assertTrue(
+        os.path.exists(os.path.join(tmp_dir, 'variables', 'variables.index')))
+    self.assertTrue(
+        os.path.exists(
+            os.path.join(tmp_dir, 'variables',
+                         'variables.data-00000-of-00001')))
 
     imported = tf.saved_model.load(tmp_dir)
     segmentation_fn = imported.signatures['serving_default']
@@ -119,9 +95,11 @@ class SemanticSegmentationExportTest(tf.test.TestCase, parameterized.TestCase):
             fn_output_signature=tf.TensorSpec(
                 shape=[112, 112, 3], dtype=tf.float32)))
     expected_output = tf.image.resize(
-        model(processed_images, training=False), [112, 112], method='bilinear')
+        module.model(processed_images, training=False), [112, 112],
+        method='bilinear')
     out = segmentation_fn(tf.constant(images))
     self.assertAllClose(out['predicted_masks'].numpy(), expected_output.numpy())
 
+
 if __name__ == '__main__':
   tf.test.main()

official/vision/beta/tasks/maskrcnn.py

@@ -341,5 +341,5 @@ class MaskRCNNTask(base_task.Task):
         step_outputs[self.coco_metric.name][1])
     return state
 
-  def reduce_aggregated_logs(self, aggregated_logs):
+  def reduce_aggregated_logs(self, aggregated_logs, global_step=None):
     return self.coco_metric.result()

official/vision/beta/tasks/retinanet.py

@@ -292,5 +292,5 @@ class RetinaNetTask(base_task.Task):
                                   step_outputs[self.coco_metric.name][1])
     return state
 
-  def reduce_aggregated_logs(self, aggregated_logs):
+  def reduce_aggregated_logs(self, aggregated_logs, global_step=None):
     return self.coco_metric.result()

official/vision/beta/tasks/semantic_segmentation.py

@@ -263,7 +263,7 @@ class SemanticSegmentationTask(base_task.Task):
                                  step_outputs[self.iou_metric.name][1])
     return state
 
-  def reduce_aggregated_logs(self, aggregated_logs):
+  def reduce_aggregated_logs(self, aggregated_logs, global_step=None):
     result = {}
     ious = self.iou_metric.result()
     # TODO(arashwan): support loading class name from a label map file.

official/vision/beta/train.py

@@ -47,7 +47,8 @@ def main(_):
   # dtype is float16
   if params.runtime.mixed_precision_dtype:
     performance.set_mixed_precision_policy(params.runtime.mixed_precision_dtype,
-                                           params.runtime.loss_scale)
+                                           params.runtime.loss_scale,
+                                           use_experimental_api=True)
   distribution_strategy = distribute_utils.get_distribution_strategy(
       distribution_strategy=params.runtime.distribution_strategy,
       all_reduce_alg=params.runtime.all_reduce_alg,

official/vision/beta/train_spatial_partitioning.py

@@ -97,7 +97,8 @@ def main(_):
   # dtype is float16
   if params.runtime.mixed_precision_dtype:
     performance.set_mixed_precision_policy(params.runtime.mixed_precision_dtype,
-                                           params.runtime.loss_scale)
+                                           params.runtime.loss_scale,
+                                           use_experimental_api=True)
 
   input_partition_dims = None
   if FLAGS.mode == 'train_and_eval':

official/vision/detection/modeling/architecture/fpn.py

@@ -28,7 +28,6 @@ import functools
 
 import tensorflow as tf
 
-from official.vision.detection.modeling.architecture import keras_utils
 from official.vision.detection.modeling.architecture import nn_ops
 from official.vision.detection.ops import spatial_transform_ops
 
@@ -120,7 +119,7 @@ class Fpn(object):
           'The minimum backbone level %d should be '%(min(input_levels)) +
           'less or equal to FPN minimum level %d.:'%(self._min_level))
     backbone_max_level = min(max(input_levels), self._max_level)
-    with keras_utils.maybe_enter_backend_graph(), tf.name_scope('fpn'):
+    with tf.name_scope('fpn'):
       # Adds lateral connections.
       feats_lateral = {}
       for level in range(self._min_level, backbone_max_level + 1):

official/vision/detection/modeling/architecture/heads.py

@@ -23,7 +23,6 @@ import functools
 import numpy as np
 import tensorflow as tf
 
-from official.vision.detection.modeling.architecture import keras_utils
 from official.vision.detection.modeling.architecture import nn_ops
 from official.vision.detection.ops import spatial_transform_ops
 
@@ -60,6 +59,8 @@ class RpnHead(tf.keras.layers.Layer):
       norm_activation: an operation that includes a normalization layer followed
         by an optional activation layer.
     """
+    super().__init__(autocast=False)
+
     self._min_level = min_level
     self._max_level = max_level
     self._anchors_per_location = anchors_per_location
@@ -123,12 +124,12 @@ class RpnHead(tf.keras.layers.Layer):
 
     return scores, bboxes
 
-  def __call__(self, features, is_training=None):
+  def call(self, features, is_training=None):
 
     scores_outputs = {}
     box_outputs = {}
 
-    with keras_utils.maybe_enter_backend_graph(), tf.name_scope('rpn_head'):
+    with tf.name_scope('rpn_head'):
       for level in range(self._min_level, self._max_level + 1):
         scores_output, box_output = self._shared_rpn_heads(
             features[level], self._anchors_per_location, level, is_training)
@@ -251,7 +252,7 @@ class OlnRpnHead(tf.keras.layers.Layer):
     box_outputs = {}
     center_outputs = {}
 
-    with keras_utils.maybe_enter_backend_graph(), tf.name_scope('rpn_head'):
+    with tf.name_scope('rpn_head'):
       for level in range(self._min_level, self._max_level + 1):
         scores_output, box_output, center_output = self._shared_rpn_heads(
             features[level], self._anchors_per_location, level, is_training)
@@ -294,6 +295,8 @@ class FastrcnnHead(tf.keras.layers.Layer):
       norm_activation: an operation that includes a normalization layer followed
         by an optional activation layer.
     """
+    super(FastrcnnHead, self).__init__(autocast=False)
+
     self._num_classes = num_classes
 
     self._num_convs = num_convs
@@ -360,7 +363,7 @@ class FastrcnnHead(tf.keras.layers.Layer):
         bias_initializer=tf.zeros_initializer(),
         name='box-predict')
 
-  def __call__(self, roi_features, is_training=None):
+  def call(self, roi_features, is_training=None):
     """Box and class branches for the Mask-RCNN model.
 
     Args:
@@ -376,7 +379,7 @@ class FastrcnnHead(tf.keras.layers.Layer):
         predictions.
     """
 
-    with keras_utils.maybe_enter_backend_graph(), tf.name_scope(
+    with tf.name_scope(
         'fast_rcnn_head'):
       # reshape inputs beofre FC.
       _, num_rois, height, width, filters = roi_features.get_shape().as_list()
@@ -520,8 +523,7 @@ class OlnBoxScoreHead(tf.keras.layers.Layer):
         predictions.
     """
 
-    with keras_utils.maybe_enter_backend_graph(), tf.name_scope(
-        'fast_rcnn_head'):
+    with tf.name_scope('fast_rcnn_head'):
       # reshape inputs beofre FC.
       _, num_rois, height, width, filters = roi_features.get_shape().as_list()
 
@@ -574,6 +576,7 @@ class MaskrcnnHead(tf.keras.layers.Layer):
       norm_activation: an operation that includes a normalization layer followed
         by an optional activation layer.
     """
+    super(MaskrcnnHead, self).__init__(autocast=False)
     self._num_classes = num_classes
     self._mask_target_size = mask_target_size
 
@@ -621,7 +624,15 @@ class MaskrcnnHead(tf.keras.layers.Layer):
         bias_initializer=tf.zeros_initializer(),
         name='conv5-mask')
 
-  def __call__(self, roi_features, class_indices, is_training=None):
+    with tf.name_scope('mask_head'):
+      self._mask_conv2d_op = self._conv2d_op(
+          self._num_classes,
+          kernel_size=(1, 1),
+          strides=(1, 1),
+          padding='valid',
+          name='mask_fcn_logits')
+
+  def call(self, roi_features, class_indices, is_training=None):
     """Mask branch for the Mask-RCNN model.
 
     Args:
@@ -642,45 +653,38 @@ class MaskrcnnHead(tf.keras.layers.Layer):
         boxes is not 4.
     """
 
-    with keras_utils.maybe_enter_backend_graph():
-      with tf.name_scope('mask_head'):
-        _, num_rois, height, width, filters = roi_features.get_shape().as_list()
-        net = tf.reshape(roi_features, [-1, height, width, filters])
-
-        for i in range(self._num_convs):
-          net = self._conv2d_ops[i](net)
-          if self._use_batch_norm:
-            net = self._norm_activation()(net, is_training=is_training)
+    with tf.name_scope('mask_head'):
+      _, num_rois, height, width, filters = roi_features.get_shape().as_list()
+      net = tf.reshape(roi_features, [-1, height, width, filters])
 
-        net = self._mask_conv_transpose(net)
+      for i in range(self._num_convs):
+        net = self._conv2d_ops[i](net)
         if self._use_batch_norm:
           net = self._norm_activation()(net, is_training=is_training)
 
-        mask_outputs = self._conv2d_op(
-            self._num_classes,
-            kernel_size=(1, 1),
-            strides=(1, 1),
-            padding='valid',
-            name='mask_fcn_logits')(
-                net)
-        mask_outputs = tf.reshape(mask_outputs, [
-            -1, num_rois, self._mask_target_size, self._mask_target_size,
-            self._num_classes
-        ])
-
-        with tf.name_scope('masks_post_processing'):
-          # TODO(pengchong): Figure out the way not to use the static inferred
-          # batch size.
-          batch_size, num_masks = class_indices.get_shape().as_list()
-          mask_outputs = tf.transpose(a=mask_outputs, perm=[0, 1, 4, 2, 3])
-          # Contructs indices for gather.
-          batch_indices = tf.tile(
-              tf.expand_dims(tf.range(batch_size), axis=1), [1, num_masks])
-          mask_indices = tf.tile(
-              tf.expand_dims(tf.range(num_masks), axis=0), [batch_size, 1])
-          gather_indices = tf.stack(
-              [batch_indices, mask_indices, class_indices], axis=2)
-          mask_outputs = tf.gather_nd(mask_outputs, gather_indices)
+      net = self._mask_conv_transpose(net)
+      if self._use_batch_norm:
+        net = self._norm_activation()(net, is_training=is_training)
+
+      mask_outputs = self._mask_conv2d_op(net)
+      mask_outputs = tf.reshape(mask_outputs, [
+          -1, num_rois, self._mask_target_size, self._mask_target_size,
+          self._num_classes
+      ])
+
+      with tf.name_scope('masks_post_processing'):
+        # TODO(pengchong): Figure out the way not to use the static inferred
+        # batch size.
+        batch_size, num_masks = class_indices.get_shape().as_list()
+        mask_outputs = tf.transpose(a=mask_outputs, perm=[0, 1, 4, 2, 3])
+        # Constructs indices for gather.
+        batch_indices = tf.tile(
+            tf.expand_dims(tf.range(batch_size), axis=1), [1, num_masks])
+        mask_indices = tf.tile(
+            tf.expand_dims(tf.range(num_masks), axis=0), [batch_size, 1])
+        gather_indices = tf.stack(
+            [batch_indices, mask_indices, class_indices], axis=2)
+        mask_outputs = tf.gather_nd(mask_outputs, gather_indices)
       return mask_outputs
 
 
@@ -826,8 +830,7 @@ class RetinanetHead(object):
     """Returns outputs of RetinaNet head."""
     class_outputs = {}
     box_outputs = {}
-    with keras_utils.maybe_enter_backend_graph(), tf.name_scope(
-        'retinanet_head'):
+    with tf.name_scope('retinanet_head'):
       for level in range(self._min_level, self._max_level + 1):
         features = fpn_features[level]
 
@@ -915,7 +918,7 @@ class ShapemaskPriorHead(object):
       detection_priors: A float Tensor of shape [batch_size * num_instances,
         mask_size, mask_size, 1].
     """
-    with keras_utils.maybe_enter_backend_graph(), tf.name_scope('prior_mask'):
+    with tf.name_scope('prior_mask'):
       batch_size, num_instances, _ = boxes.get_shape().as_list()
       outer_boxes = tf.cast(outer_boxes, tf.float32)
       boxes = tf.cast(boxes, tf.float32)
@@ -991,8 +994,8 @@ class ShapemaskPriorHead(object):
     features = tf.reduce_mean(features, axis=(2, 3))
     logits = tf.keras.layers.Dense(
         self._mask_num_classes * self._num_clusters,
-        kernel_initializer=tf.random_normal_initializer(stddev=0.01))(
-            features)
+        kernel_initializer=tf.random_normal_initializer(stddev=0.01),
+        name='classify-shape-prior-fc')(features)
     logits = tf.reshape(
         logits,
         [batch_size, num_instances, self._mask_num_classes, self._num_clusters])
@@ -1085,7 +1088,7 @@ class ShapemaskCoarsemaskHead(object):
       mask_outputs: instance mask prediction as a float Tensor of shape
         [batch_size, num_instances, mask_size, mask_size].
     """
-    with keras_utils.maybe_enter_backend_graph(), tf.name_scope('coarse_mask'):
+    with tf.name_scope('coarse_mask'):
       # Transform detection priors to have the same dimension as features.
       detection_priors = tf.expand_dims(detection_priors, axis=-1)
       detection_priors = self._coarse_mask_fc(detection_priors)
@@ -1217,7 +1220,7 @@ class ShapemaskFinemaskHead(object):
     """
     # Extract the foreground mean features
     # with tf.variable_scope('fine_mask', reuse=tf.AUTO_REUSE):
-    with keras_utils.maybe_enter_backend_graph(), tf.name_scope('fine_mask'):
+    with tf.name_scope('fine_mask'):
       mask_probs = tf.nn.sigmoid(mask_logits)
       # Compute instance embedding for hard average.
       binary_mask = tf.cast(tf.greater(mask_probs, 0.5), features.dtype)

official/vision/detection/modeling/architecture/resnet.py

@@ -23,9 +23,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from absl import logging
 import tensorflow as tf
-from official.vision.detection.modeling.architecture import keras_utils
 from official.vision.detection.modeling.architecture import nn_ops
 
 
@@ -112,9 +110,8 @@ class Resnet(object):
       The values are corresponding feature hierarchy in ResNet with shape
       [batch_size, height_l, width_l, num_filters].
     """
-    with keras_utils.maybe_enter_backend_graph():
-      with tf.name_scope('resnet%s' % self._resnet_depth):
-        return self._resnet_fn(inputs, is_training)
+    with tf.name_scope('resnet%s' % self._resnet_depth):
+      return self._resnet_fn(inputs, is_training)
 
   def fixed_padding(self, inputs, kernel_size):
     """Pads the input along the spatial dimensions independently of input size.

official/vision/detection/modeling/architecture/spinenet.py

@@ -25,7 +25,6 @@ from absl import logging
 import tensorflow as tf
 
 from official.modeling import tf_utils
-from official.vision.detection.modeling.architecture import keras_utils
 from official.vision.detection.modeling.architecture import nn_blocks
 
 layers = tf.keras.layers
@@ -486,21 +485,20 @@ class SpineNetBuilder(object):
     self._norm_epsilon = norm_epsilon
 
   def __call__(self, inputs, is_training=None):
-    with keras_utils.maybe_enter_backend_graph():
-      model = SpineNet(
-          input_specs=self._input_specs,
-          min_level=self._min_level,
-          max_level=self._max_level,
-          block_specs=self._block_specs,
-          endpoints_num_filters=self._endpoints_num_filters,
-          resample_alpha=self._resample_alpha,
-          block_repeats=self._block_repeats,
-          filter_size_scale=self._filter_size_scale,
-          kernel_initializer=self._kernel_initializer,
-          kernel_regularizer=self._kernel_regularizer,
-          bias_regularizer=self._bias_regularizer,
-          activation=self._activation,
-          use_sync_bn=self._use_sync_bn,
-          norm_momentum=self._norm_momentum,
-          norm_epsilon=self._norm_epsilon)
-      return model(inputs)
+    model = SpineNet(
+        input_specs=self._input_specs,
+        min_level=self._min_level,
+        max_level=self._max_level,
+        block_specs=self._block_specs,
+        endpoints_num_filters=self._endpoints_num_filters,
+        resample_alpha=self._resample_alpha,
+        block_repeats=self._block_repeats,
+        filter_size_scale=self._filter_size_scale,
+        kernel_initializer=self._kernel_initializer,
+        kernel_regularizer=self._kernel_regularizer,
+        bias_regularizer=self._bias_regularizer,
+        activation=self._activation,
+        use_sync_bn=self._use_sync_bn,
+        norm_momentum=self._norm_momentum,
+        norm_epsilon=self._norm_epsilon)
+    return model(inputs)

official/vision/detection/modeling/checkpoint_utils.py

@@ -51,13 +51,16 @@ def _build_assignment_map(keras_model,
   """
   assignment_map = {}
 
-  checkpoint_names = None
+  checkpoint_names = []
   if var_to_shape_map:
     checkpoint_names = list(
         filter(
             lambda x: not x.endswith('Momentum') and not x.endswith(
                 'global_step'), var_to_shape_map.keys()))
 
+  logging.info('Number of variables in the checkpoint %d',
+               len(checkpoint_names))
+
   for var in keras_model.variables:
     var_name = var.name
 
@@ -87,7 +90,7 @@ def _build_assignment_map(keras_model,
       logging.info('Error removing the match_name: %s', match_names)
       logging.info('Exception: %s', e)
       raise
-  logging.info('Found variable in checkpoint: %d', len(assignment_map))
+  logging.info('Found matching variable in checkpoint: %d', len(assignment_map))
   return assignment_map
 
 

official/vision/detection/modeling/maskrcnn_model.py

@@ -26,7 +26,6 @@ from official.vision.detection.evaluation import factory as eval_factory
 from official.vision.detection.modeling import base_model
 from official.vision.detection.modeling import losses
 from official.vision.detection.modeling.architecture import factory
-from official.vision.detection.modeling.architecture import keras_utils
 from official.vision.detection.ops import postprocess_ops
 from official.vision.detection.ops import roi_ops
 from official.vision.detection.ops import spatial_transform_ops
@@ -293,14 +292,13 @@ class MaskrcnnModel(base_model.Model):
   def build_model(self, params, mode):
     if self._keras_model is None:
       input_layers = self.build_input_layers(self._params, mode)
-      with keras_utils.maybe_enter_backend_graph():
-        outputs = self.model_outputs(input_layers, mode)
-
-        model = tf.keras.models.Model(
-            inputs=input_layers, outputs=outputs, name='maskrcnn')
-        assert model is not None, 'Fail to build tf.keras.Model.'
-        model.optimizer = self.build_optimizer()
-        self._keras_model = model
+      outputs = self.model_outputs(input_layers, mode)
+
+      model = tf.keras.models.Model(
+          inputs=input_layers, outputs=outputs, name='maskrcnn')
+      assert model is not None, 'Fail to build tf.keras.Model.'
+      model.optimizer = self.build_optimizer()
+      self._keras_model = model
 
     return self._keras_model
 

official/vision/detection/modeling/olnmask_model.py

@@ -24,7 +24,6 @@ from official.vision.detection.dataloader import anchor
 from official.vision.detection.dataloader import mode_keys
 from official.vision.detection.modeling import losses
 from official.vision.detection.modeling.architecture import factory
-from official.vision.detection.modeling.architecture import keras_utils
 from official.vision.detection.modeling.maskrcnn_model import MaskrcnnModel
 from official.vision.detection.ops import postprocess_ops
 from official.vision.detection.ops import roi_ops
@@ -422,13 +421,12 @@ class OlnMaskModel(MaskrcnnModel):
   def build_model(self, params, mode):
     if self._keras_model is None:
       input_layers = self.build_input_layers(self._params, mode)
-      with keras_utils.maybe_enter_backend_graph():
-        outputs = self.model_outputs(input_layers, mode)
-
-        model = tf.keras.models.Model(
-            inputs=input_layers, outputs=outputs, name='olnmask')
-        assert model is not None, 'Fail to build tf.keras.Model.'
-        model.optimizer = self.build_optimizer()
-        self._keras_model = model
+      outputs = self.model_outputs(input_layers, mode)
+
+      model = tf.keras.models.Model(
+          inputs=input_layers, outputs=outputs, name='olnmask')
+      assert model is not None, 'Fail to build tf.keras.Model.'
+      model.optimizer = self.build_optimizer()
+      self._keras_model = model
 
     return self._keras_model

official/vision/detection/modeling/retinanet_model.py

@@ -25,7 +25,6 @@ from official.vision.detection.evaluation import factory as eval_factory
 from official.vision.detection.modeling import base_model
 from official.vision.detection.modeling import losses
 from official.vision.detection.modeling.architecture import factory
-from official.vision.detection.modeling.architecture import keras_utils
 from official.vision.detection.ops import postprocess_ops
 
 
@@ -117,14 +116,13 @@ class RetinanetModel(base_model.Model):
 
   def build_model(self, params, mode=None):
     if self._keras_model is None:
-      with keras_utils.maybe_enter_backend_graph():
-        outputs = self.model_outputs(self._input_layer, mode)
-
-        model = tf.keras.models.Model(
-            inputs=self._input_layer, outputs=outputs, name='retinanet')
-        assert model is not None, 'Fail to build tf.keras.Model.'
-        model.optimizer = self.build_optimizer()
-        self._keras_model = model
+      outputs = self.model_outputs(self._input_layer, mode)
+
+      model = tf.keras.models.Model(
+          inputs=self._input_layer, outputs=outputs, name='retinanet')
+      assert model is not None, 'Fail to build tf.keras.Model.'
+      model.optimizer = self.build_optimizer()
+      self._keras_model = model
 
     return self._keras_model
 

official/vision/detection/modeling/shapemask_model.py

@@ -26,7 +26,6 @@ from official.vision.detection.evaluation import factory as eval_factory
 from official.vision.detection.modeling import base_model
 from official.vision.detection.modeling import losses
 from official.vision.detection.modeling.architecture import factory
-from official.vision.detection.modeling.architecture import keras_utils
 from official.vision.detection.ops import postprocess_ops
 from official.vision.detection.utils import box_utils
 
@@ -256,14 +255,13 @@ class ShapeMaskModel(base_model.Model):
   def build_model(self, params, mode):
     if self._keras_model is None:
       input_layers = self.build_input_layers(self._params, mode)
-      with keras_utils.maybe_enter_backend_graph():
-        outputs = self.model_outputs(input_layers, mode)
-
-        model = tf.keras.models.Model(
-            inputs=input_layers, outputs=outputs, name='shapemask')
-        assert model is not None, 'Fail to build tf.keras.Model.'
-        model.optimizer = self.build_optimizer()
-        self._keras_model = model
+      outputs = self.model_outputs(input_layers, mode)
+
+      model = tf.keras.models.Model(
+          inputs=input_layers, outputs=outputs, name='shapemask')
+      assert model is not None, 'Fail to build tf.keras.Model.'
+      model.optimizer = self.build_optimizer()
+      self._keras_model = model
 
     return self._keras_model
 

official/vision/detection/ops/postprocess_ops.py

@@ -291,15 +291,16 @@ def _generate_detections_batched(boxes, scores, max_total_size,
   return nmsed_boxes, nmsed_scores, nmsed_classes, valid_detections
 
 
-class MultilevelDetectionGenerator(object):
+class MultilevelDetectionGenerator(tf.keras.layers.Layer):
   """Generates detected boxes with scores and classes for one-stage detector."""
 
   def __init__(self, min_level, max_level, params):
     self._min_level = min_level
     self._max_level = max_level
     self._generate_detections = generate_detections_factory(params)
+    super(MultilevelDetectionGenerator, self).__init__(autocast=False)
 
-  def __call__(self, box_outputs, class_outputs, anchor_boxes, image_shape):
+  def call(self, box_outputs, class_outputs, anchor_boxes, image_shape):
     # Collects outputs from all levels into a list.
     boxes = []
     scores = []
@@ -337,13 +338,14 @@ class MultilevelDetectionGenerator(object):
     return nmsed_boxes, nmsed_scores, nmsed_classes, valid_detections
 
 
-class GenericDetectionGenerator(object):
+class GenericDetectionGenerator(tf.keras.layers.Layer):
   """Generates the final detected boxes with scores and classes."""
 
   def __init__(self, params):
+    super(GenericDetectionGenerator, self).__init__(autocast=False)
     self._generate_detections = generate_detections_factory(params)
 
-  def __call__(self, box_outputs, class_outputs, anchor_boxes, image_shape):
+  def call(self, box_outputs, class_outputs, anchor_boxes, image_shape):
     """Generate final detections.
 
     Args:

official/vision/detection/ops/roi_ops.py

@@ -170,7 +170,7 @@ def multilevel_propose_rois(rpn_boxes,
     return selected_rois, selected_roi_scores
 
 
-class ROIGenerator(object):
+class ROIGenerator(tf.keras.layers.Layer):
   """Proposes RoIs for the second stage processing."""
 
   def __init__(self, params):
@@ -185,8 +185,9 @@ class ROIGenerator(object):
     self._test_rpn_score_threshold = params.test_rpn_score_threshold
     self._test_rpn_min_size_threshold = params.test_rpn_min_size_threshold
     self._use_batched_nms = params.use_batched_nms
+    super(ROIGenerator, self).__init__(autocast=False)
 
-  def __call__(self, boxes, scores, anchor_boxes, image_shape, is_training):
+  def call(self, boxes, scores, anchor_boxes, image_shape, is_training):
     """Generates RoI proposals.
 
     Args:

official/vision/detection/ops/target_ops.py

@@ -292,7 +292,7 @@ def sample_and_crop_foreground_masks(candidate_rois,
     return foreground_rois, foreground_classes, cropped_foreground_masks
 
 
-class ROISampler(object):
+class ROISampler(tf.keras.layers.Layer):
   """Samples RoIs and creates training targets."""
 
   def __init__(self, params):
@@ -302,8 +302,9 @@ class ROISampler(object):
     self._bg_iou_thresh_hi = params.bg_iou_thresh_hi
     self._bg_iou_thresh_lo = params.bg_iou_thresh_lo
     self._mix_gt_boxes = params.mix_gt_boxes
+    super(ROISampler, self).__init__(autocast=False)
 
-  def __call__(self, rois, gt_boxes, gt_classes):
+  def call(self, rois, gt_boxes, gt_classes):
     """Sample and assign RoIs for training.
 
     Args:
@@ -516,15 +517,20 @@ class ROIScoreSampler(ROISampler):
               sampled_gt_classes, sampled_gt_indices)
 
 
-class MaskSampler(object):
+class MaskSampler(tf.keras.layers.Layer):
   """Samples and creates mask training targets."""
 
   def __init__(self, mask_target_size, num_mask_samples_per_image):
     self._mask_target_size = mask_target_size
     self._num_mask_samples_per_image = num_mask_samples_per_image
-
-  def __call__(self, candidate_rois, candidate_gt_boxes, candidate_gt_classes,
-               candidate_gt_indices, gt_masks):
+    super(MaskSampler, self).__init__(autocast=False)
+
+  def call(self,
+           candidate_rois,
+           candidate_gt_boxes,
+           candidate_gt_classes,
+           candidate_gt_indices,
+           gt_masks):
     """Sample and create mask targets for training.
 
     Args:

official/vision/image_classification/classifier_trainer.py

@@ -341,7 +341,8 @@ def train_and_eval(
     optimizer = performance.configure_optimizer(
         optimizer,
         use_float16=train_builder.dtype == 'float16',
-        loss_scale=get_loss_scale(params))
+        loss_scale=get_loss_scale(params),
+        use_experimental_api=True)
 
     metrics_map = _get_metrics(one_hot)
     metrics = [metrics_map[metric] for metric in params.train.metrics]

research/deeplab/README.md

@@ -246,7 +246,7 @@ PASCAL VOC 2012 and Cityscapes.
 ### March 5, 2018
 
 * First release of DeepLab in TensorFlow including deeper Xception network
-backbone. Included chekcpoints that have been pretrained on PASCAL VOC 2012
+backbone. Included checkpoints that have been pretrained on PASCAL VOC 2012
 and Cityscapes.
 
 ## References