Internal change

PiperOrigin-RevId: 361196060

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/keras_utils.py

-# Copyright 2019 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.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Util functions to integrate with Keras internals."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-
-from tensorflow.python.keras import backend
-
-try:
-  from tensorflow.python.keras.engine import keras_tensor  # pylint: disable=g-import-not-at-top,unused-import
-  keras_tensor.disable_keras_tensors()
-except ImportError:
-  keras_tensor = None
-
-
-class NoOpContextManager(object):
-
-  def __enter__(self):
-    pass
-
-  def __exit__(self, *args):
-    pass
-
-
-def maybe_enter_backend_graph():
-  if (keras_tensor is not None) and keras_tensor.keras_tensors_enabled():
-    return NoOpContextManager()
-  else:
-    return backend.get_graph().as_default()

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: