Merge branch 'master' into absl

official/vision/detection/modeling/maskrcnn_model.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.
+# ==============================================================================
+"""Model defination for the Mask R-CNN Model."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow.compat.v2 as tf
+
+from tensorflow.python.keras import backend
+from official.vision.detection.dataloader import anchor
+from official.vision.detection.dataloader import mode_keys
+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.ops import postprocess_ops
+from official.vision.detection.ops import roi_ops
+from official.vision.detection.ops import sampling_ops
+from official.vision.detection.ops import spatial_transform_ops
+from official.vision.detection.utils import box_utils
+
+
+class MaskrcnnModel(base_model.Model):
+  """Mask R-CNN model function."""
+
+  def __init__(self, params):
+    super(MaskrcnnModel, self).__init__(params)
+
+    # For eval metrics.
+    self._params = params
+    self._keras_model = None
+
+    self._include_mask = params.architecture.include_mask
+
+    # Architecture generators.
+    self._backbone_fn = factory.backbone_generator(params)
+    self._fpn_fn = factory.multilevel_features_generator(params)
+    self._rpn_head_fn = factory.rpn_head_generator(params.rpn_head)
+    self._generate_rois_fn = roi_ops.ROIGenerator(params.roi_proposal)
+    self._sample_rois_fn = sampling_ops.ROISampler(params.roi_sampling)
+    self._sample_masks_fn = sampling_ops.MaskSampler(params.mask_sampling)
+
+    self._frcnn_head_fn = factory.fast_rcnn_head_generator(params.frcnn_head)
+    if self._include_mask:
+      self._mrcnn_head_fn = factory.mask_rcnn_head_generator(params.mrcnn_head)
+
+    # Loss function.
+    self._rpn_score_loss_fn = losses.RpnScoreLoss(params.rpn_score_loss)
+    self._rpn_box_loss_fn = losses.RpnBoxLoss(params.rpn_box_loss)
+    self._frcnn_class_loss_fn = losses.FastrcnnClassLoss()
+    self._frcnn_box_loss_fn = losses.FastrcnnBoxLoss(params.frcnn_box_loss)
+    if self._include_mask:
+      self._mask_loss_fn = losses.MaskrcnnLoss()
+
+    self._generate_detections_fn = postprocess_ops.GenericDetectionGenerator(
+        params.postprocess)
+
+    self._transpose_input = params.train.transpose_input
+    assert not self._transpose_input, 'Transpose input is not supportted.'
+
+  def build_outputs(self, inputs, mode):
+    is_training = mode == mode_keys.TRAIN
+    model_outputs = {}
+
+    image = inputs['image']
+    _, image_height, image_width, _ = image.get_shape().as_list()
+    backbone_features = self._backbone_fn(image, is_training)
+    fpn_features = self._fpn_fn(backbone_features, is_training)
+
+    rpn_score_outputs, rpn_box_outputs = self._rpn_head_fn(
+        fpn_features, is_training)
+    model_outputs.update({
+        'rpn_score_outputs':
+            tf.nest.map_structure(lambda x: tf.cast(x, tf.float32),
+                                  rpn_score_outputs),
+        'rpn_box_outputs':
+            tf.nest.map_structure(lambda x: tf.cast(x, tf.float32),
+                                  rpn_box_outputs),
+    })
+    input_anchor = anchor.Anchor(self._params.anchor.min_level,
+                                 self._params.anchor.max_level,
+                                 self._params.anchor.num_scales,
+                                 self._params.anchor.aspect_ratios,
+                                 self._params.anchor.anchor_size,
+                                 (image_height, image_width))
+    rpn_rois, _ = self._generate_rois_fn(rpn_box_outputs, rpn_score_outputs,
+                                         input_anchor.multilevel_boxes,
+                                         inputs['image_info'][:, 1, :],
+                                         is_training)
+    if is_training:
+      rpn_rois = tf.stop_gradient(rpn_rois)
+
+      # Sample proposals.
+      rpn_rois, matched_gt_boxes, matched_gt_classes, matched_gt_indices = (
+          self._sample_rois_fn(rpn_rois, inputs['gt_boxes'],
+                               inputs['gt_classes']))
+
+      # Create bounding box training targets.
+      box_targets = box_utils.encode_boxes(
+          matched_gt_boxes, rpn_rois, weights=[10.0, 10.0, 5.0, 5.0])
+      # If the target is background, the box target is set to all 0s.
+      box_targets = tf.where(
+          tf.tile(
+              tf.expand_dims(tf.equal(matched_gt_classes, 0), axis=-1),
+              [1, 1, 4]),
+          tf.zeros_like(box_targets),
+          box_targets)
+      model_outputs.update({
+          'class_targets': matched_gt_classes,
+          'box_targets': box_targets,
+      })
+
+    roi_features = spatial_transform_ops.multilevel_crop_and_resize(
+        fpn_features, rpn_rois, output_size=7)
+
+    class_outputs, box_outputs = self._frcnn_head_fn(roi_features, is_training)
+
+    model_outputs.update({
+        'class_outputs':
+            tf.nest.map_structure(lambda x: tf.cast(x, tf.float32),
+                                  class_outputs),
+        'box_outputs':
+            tf.nest.map_structure(lambda x: tf.cast(x, tf.float32),
+                                  box_outputs),
+    })
+
+    # Add this output to train to make the checkpoint loadable in predict mode.
+    # If we skip it in train mode, the heads will be out-of-order and checkpoint
+    # loading will fail.
+    boxes, scores, classes, valid_detections = self._generate_detections_fn(
+        box_outputs, class_outputs, rpn_rois, inputs['image_info'][:, 1:2, :])
+    model_outputs.update({
+        'num_detections': valid_detections,
+        'detection_boxes': boxes,
+        'detection_classes': classes,
+        'detection_scores': scores,
+    })
+
+    if not self._include_mask:
+      return model_outputs
+
+    if is_training:
+      rpn_rois, classes, mask_targets = self._sample_masks_fn(
+          rpn_rois, matched_gt_boxes, matched_gt_classes, matched_gt_indices,
+          inputs['gt_masks'])
+      mask_targets = tf.stop_gradient(mask_targets)
+
+      classes = tf.cast(classes, dtype=tf.int32)
+
+      model_outputs.update({
+          'mask_targets': mask_targets,
+          'sampled_class_targets': classes,
+      })
+    else:
+      rpn_rois = boxes
+      classes = tf.cast(classes, dtype=tf.int32)
+
+    mask_roi_features = spatial_transform_ops.multilevel_crop_and_resize(
+        fpn_features, rpn_rois, output_size=14)
+
+    mask_outputs = self._mrcnn_head_fn(mask_roi_features, classes, is_training)
+
+    if is_training:
+      model_outputs.update({
+          'mask_outputs':
+              tf.nest.map_structure(lambda x: tf.cast(x, tf.float32),
+                                    mask_outputs),
+      })
+    else:
+      model_outputs.update({
+          'detection_masks': tf.nn.sigmoid(mask_outputs)
+      })
+
+    return model_outputs
+
+  def build_loss_fn(self):
+    if self._keras_model is None:
+      raise ValueError('build_loss_fn() must be called after build_model().')
+
+    filter_fn = self.make_filter_trainable_variables_fn()
+    trainable_variables = filter_fn(self._keras_model.trainable_variables)
+
+    def _total_loss_fn(labels, outputs):
+      rpn_score_loss = self._rpn_score_loss_fn(outputs['rpn_score_outputs'],
+                                               labels['rpn_score_targets'])
+      rpn_box_loss = self._rpn_box_loss_fn(outputs['rpn_box_outputs'],
+                                           labels['rpn_box_targets'])
+
+      frcnn_class_loss = self._frcnn_class_loss_fn(outputs['class_outputs'],
+                                                   outputs['class_targets'])
+      frcnn_box_loss = self._frcnn_box_loss_fn(outputs['box_outputs'],
+                                               outputs['class_targets'],
+                                               outputs['box_targets'])
+
+      if self._include_mask:
+        mask_loss = self._mask_loss_fn(outputs['mask_outputs'],
+                                       outputs['mask_targets'],
+                                       outputs['sampled_class_targets'])
+      else:
+        mask_loss = 0.0
+
+      model_loss = (
+          rpn_score_loss + rpn_box_loss + frcnn_class_loss + frcnn_box_loss +
+          mask_loss)
+
+      l2_regularization_loss = self.weight_decay_loss(trainable_variables)
+      total_loss = model_loss + l2_regularization_loss
+      return {
+          'total_loss': total_loss,
+          'loss': total_loss,
+          'fast_rcnn_class_loss': frcnn_class_loss,
+          'fast_rcnn_box_loss': frcnn_box_loss,
+          'mask_loss': mask_loss,
+          'model_loss': model_loss,
+          'l2_regularization_loss': l2_regularization_loss,
+          'rpn_score_loss': rpn_score_loss,
+          'rpn_box_loss': rpn_box_loss,
+      }
+
+    return _total_loss_fn
+
+  def build_input_layers(self, params, mode):
+    is_training = mode == mode_keys.TRAIN
+    input_shape = (
+        params.maskrcnn_parser.output_size +
+        [params.maskrcnn_parser.num_channels])
+    if is_training:
+      batch_size = params.train.batch_size
+      input_layer = {
+          'image':
+              tf.keras.layers.Input(
+                  shape=input_shape,
+                  batch_size=batch_size,
+                  name='image',
+                  dtype=tf.bfloat16 if self._use_bfloat16 else tf.float32),
+          'image_info':
+              tf.keras.layers.Input(
+                  shape=[4, 2],
+                  batch_size=batch_size,
+                  name='image_info',
+              ),
+          'gt_boxes':
+              tf.keras.layers.Input(
+                  shape=[params.maskrcnn_parser.max_num_instances, 4],
+                  batch_size=batch_size,
+                  name='gt_boxes'),
+          'gt_classes':
+              tf.keras.layers.Input(
+                  shape=[params.maskrcnn_parser.max_num_instances],
+                  batch_size=batch_size,
+                  name='gt_classes',
+                  dtype=tf.int64),
+      }
+      if self._include_mask:
+        input_layer['gt_masks'] = tf.keras.layers.Input(
+            shape=[
+                params.maskrcnn_parser.max_num_instances,
+                params.maskrcnn_parser.mask_crop_size,
+                params.maskrcnn_parser.mask_crop_size
+            ],
+            batch_size=batch_size,
+            name='gt_masks')
+    else:
+      batch_size = params.eval.batch_size
+      input_layer = {
+          'image':
+              tf.keras.layers.Input(
+                  shape=input_shape,
+                  batch_size=batch_size,
+                  name='image',
+                  dtype=tf.bfloat16 if self._use_bfloat16 else tf.float32),
+          'image_info':
+              tf.keras.layers.Input(
+                  shape=[4, 2],
+                  batch_size=batch_size,
+                  name='image_info',
+              ),
+      }
+    return input_layer
+
+  def build_model(self, params, mode):
+    if self._keras_model is None:
+      input_layers = self.build_input_layers(self._params, mode)
+      with backend.get_graph().as_default():
+        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
+
+  def post_processing(self, labels, outputs):
+    required_output_fields = ['class_outputs', 'box_outputs']
+    for field in required_output_fields:
+      if field not in outputs:
+        raise ValueError('"%s" is missing in outputs, requried %s found %s'
+                         %(field, required_output_fields, outputs.keys()))
+    predictions = {
+        'image_info': labels['image_info'],
+        'num_detections': outputs['num_detections'],
+        'detection_boxes': outputs['detection_boxes'],
+        'detection_classes': outputs['detection_classes'],
+        'detection_scores': outputs['detection_scores'],
+    }
+    if self._include_mask:
+      predictions.update({
+          'detection_masks': outputs['detection_masks'],
+      })
+
+    if 'groundtruths' in labels:
+      predictions['source_id'] = labels['groundtruths']['source_id']
+      predictions['gt_source_id'] = labels['groundtruths']['source_id']
+      predictions['gt_height'] = labels['groundtruths']['height']
+      predictions['gt_width'] = labels['groundtruths']['width']
+      predictions['gt_image_info'] = labels['image_info']
+      predictions['gt_num_detections'] = (
+          labels['groundtruths']['num_detections'])
+      predictions['gt_boxes'] = labels['groundtruths']['boxes']
+      predictions['gt_classes'] = labels['groundtruths']['classes']
+      predictions['gt_areas'] = labels['groundtruths']['areas']
+      predictions['gt_is_crowds'] = labels['groundtruths']['is_crowds']
+    return labels, predictions
+
+  def eval_metrics(self):
+    return eval_factory.evaluator_generator(self._params.eval)

official/vision/detection/utils/object_detection/visualization_utils.py

@@ -21,6 +21,7 @@ The functions do not return a value, instead they modify the image itself.
 """
 import collections
 import functools
+from absl import logging
 # Set headless-friendly backend.
 import matplotlib; matplotlib.use('Agg')  # pylint: disable=multiple-statements
 import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
@@ -97,6 +98,12 @@ def encode_image_array_as_png_str(image):
 def visualize_images_with_bounding_boxes(images, box_outputs, step,
                                          summary_writer):
   """Records subset of evaluation images with bounding boxes."""
+  if not isinstance(images, list):
+    logging.warning('visualize_images_with_bounding_boxes expects list of '
+                    'images but received type: %s and value: %s',
+                    type(images), images)
+    return
+
   image_shape = tf.shape(images[0])
   image_height = tf.cast(image_shape[0], tf.float32)
   image_width = tf.cast(image_shape[1], tf.float32)

official/vision/image_classification/README.md

@@ -13,7 +13,7 @@ For more information about other types of models, please refer to this
 Similar to the [estimator implementation](../../r1/resnet), the Keras
 implementation has code for the ImageNet dataset. The ImageNet
 version uses a ResNet50 model implemented in
-[`resnet_model.py`](./resnet_model.py).
+[`resnet_model.py`](./resnet/resnet_model.py).
 
 Please make sure that you have the latest version of TensorFlow
 installed and
@@ -36,14 +36,14 @@ provide a few options.
 Once your dataset is ready, you can begin training the model as follows:
 
 ```bash
-python resnet_imagenet_main.py
+python resnet/resnet_imagenet_main.py
 ```
 
 Again, if you did not download the data to the default directory, specify the
 location with the `--data_dir` flag:
 
 ```bash
-python resnet_imagenet_main.py --data_dir=/path/to/imagenet
+python resnet/resnet_imagenet_main.py --data_dir=/path/to/imagenet
 ```
 
 There are more flag options you can specify. Here are some examples:
@@ -62,7 +62,7 @@ For example, this is a typical command line to run with ImageNet data with
 batch size 128 per GPU:
 
 ```bash
-python -m resnet_imagenet_main \
+python -m resnet/resnet_imagenet_main.py \
     --model_dir=/tmp/model_dir/something \
     --num_gpus=2 \
     --batch_size=128 \
@@ -120,7 +120,7 @@ From a GCE VM, you can run the following command to train ResNet for one epoch
 on a v2-8 or v3-8 TPU:
 
 ```bash
-python resnet_ctl_imagenet_main.py \
+python resnet/resnet_ctl_imagenet_main.py \
   --tpu=$TPU_NAME \
   --model_dir=$MODEL_DIR \
   --data_dir=$DATA_DIR \
@@ -140,7 +140,7 @@ python resnet_ctl_imagenet_main.py \
 To train the ResNet to convergence, run it for 90 epochs:
 
 ```bash
-python resnet_ctl_imagenet_main.py \
+python resnet/resnet_ctl_imagenet_main.py \
   --tpu=$TPU_NAME \
   --model_dir=$MODEL_DIR \
   --data_dir=$DATA_DIR \

official/vision/image_classification/common_test.py

-# Copyright 2018 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.
-# ==============================================================================
-"""Tests for the common module."""
-from __future__ import absolute_import
-from __future__ import print_function
-
-# pylint: disable=g-bad-import-order
-from mock import Mock
-import numpy as np
-import tensorflow as tf
-
-from tensorflow.python.platform import googletest
-from official.utils.misc import keras_utils
-from official.vision.image_classification import common
-
-
-class KerasCommonTests(tf.test.TestCase):
-  """Tests for common."""
-
-  @classmethod
-  def setUpClass(cls):  # pylint: disable=invalid-name
-    super(KerasCommonTests, cls).setUpClass()
-
-  def test_build_stats(self):
-
-    history = self._build_history(1.145, cat_accuracy=.99988)
-    eval_output = self._build_eval_output(.56432111, 5.990)
-    th = keras_utils.TimeHistory(128, 100)
-
-    th.timestamp_log = [keras_utils.BatchTimestamp(0, 1),
-                        keras_utils.BatchTimestamp(1, 2),
-                        keras_utils.BatchTimestamp(2, 3)]
-    th.train_finish_time = 12345
-    stats = common.build_stats(history, eval_output, [th])
-
-    self.assertEqual(1.145, stats['loss'])
-    self.assertEqual(.99988, stats['training_accuracy_top_1'])
-
-    self.assertEqual(.56432111, stats['accuracy_top_1'])
-    self.assertEqual(5.990, stats['eval_loss'])
-
-    self.assertEqual(3, stats['step_timestamp_log'][2].timestamp)
-    self.assertEqual(12345, stats['train_finish_time'])
-
-  def test_build_stats_sparse(self):
-
-    history = self._build_history(1.145, cat_accuracy_sparse=.99988)
-    eval_output = self._build_eval_output(.928, 1.9844)
-    stats = common.build_stats(history, eval_output, None)
-
-    self.assertEqual(1.145, stats['loss'])
-    self.assertEqual(.99988, stats['training_accuracy_top_1'])
-
-    self.assertEqual(.928, stats['accuracy_top_1'])
-    self.assertEqual(1.9844, stats['eval_loss'])
-
-  def test_time_history(self):
-    th = keras_utils.TimeHistory(batch_size=128, log_steps=3)
-
-    th.on_train_begin()
-    th.on_batch_begin(0)
-    th.on_batch_end(0)
-    th.on_batch_begin(1)
-    th.on_batch_end(1)
-    th.on_batch_begin(2)
-    th.on_batch_end(2)
-    th.on_batch_begin(3)
-    th.on_batch_end(3)
-    th.on_batch_begin(4)
-    th.on_batch_end(4)
-    th.on_batch_begin(5)
-    th.on_batch_end(5)
-    th.on_batch_begin(6)
-    th.on_batch_end(6)
-    th.on_train_end()
-
-    self.assertEqual(3, len(th.timestamp_log))
-
-  def _build_history(self, loss, cat_accuracy=None,
-                     cat_accuracy_sparse=None):
-    history_p = Mock()
-    history = {}
-    history_p.history = history
-    history['loss'] = [np.float64(loss)]
-    if cat_accuracy:
-      history['categorical_accuracy'] = [np.float64(cat_accuracy)]
-    if cat_accuracy_sparse:
-      history['sparse_categorical_accuracy'] = [np.float64(cat_accuracy_sparse)]
-
-    return history_p
-
-  def _build_eval_output(self, top_1, eval_loss):
-    eval_output = [np.float64(eval_loss), np.float64(top_1)]
-    return eval_output
-
-if __name__ == '__main__':
-  tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
-  googletest.main()

official/vision/image_classification/mnist_main.py

@@ -28,7 +28,7 @@ import tensorflow_datasets as tfds
 from official.utils.flags import core as flags_core
 from official.utils.misc import distribution_utils
 from official.utils.misc import model_helpers
-from official.vision.image_classification import common
+from official.vision.image_classification.resnet import common
 
 FLAGS = flags.FLAGS
 

official/vision/image_classification/cifar_preprocessing.py → official/vision/image_classification/resnet/cifar_preprocessing.py

@@ -22,7 +22,7 @@ import os
 from absl import logging
 import tensorflow as tf
 
-from official.vision.image_classification import imagenet_preprocessing
+from official.vision.image_classification.resnet import imagenet_preprocessing
 
 HEIGHT = 32
 WIDTH = 32

official/vision/image_classification/resnet_ctl_imagenet_main.py → official/vision/image_classification/resnet/resnet_ctl_imagenet_main.py

@@ -30,9 +30,9 @@ from official.utils.logs import logger
 from official.utils.misc import distribution_utils
 from official.utils.misc import keras_utils
 from official.utils.misc import model_helpers
-from official.vision.image_classification import common
-from official.vision.image_classification import imagenet_preprocessing
-from official.vision.image_classification import resnet_runnable
+from official.vision.image_classification.resnet import common
+from official.vision.image_classification.resnet import imagenet_preprocessing
+from official.vision.image_classification.resnet import resnet_runnable
 
 flags.DEFINE_boolean(name='use_tf_function', default=True,
                      help='Wrap the train and test step inside a '
@@ -147,9 +147,7 @@ def run(flags_obj):
     runnable = resnet_runnable.ResnetRunnable(flags_obj, time_callback,
                                               per_epoch_steps)
 
-  eval_interval = (
-      flags_obj.epochs_between_evals *
-      per_epoch_steps if not flags_obj.skip_eval else None)
+  eval_interval = flags_obj.epochs_between_evals * per_epoch_steps
   checkpoint_interval = (
       per_epoch_steps if flags_obj.enable_checkpoint_and_export else None)
   summary_interval = per_epoch_steps if flags_obj.enable_tensorboard else None
@@ -174,7 +172,7 @@ def run(flags_obj):
       eval_interval=eval_interval)
 
   time_callback.on_train_begin()
-  resnet_controller.train(evaluate=True)
+  resnet_controller.train(evaluate=not flags_obj.skip_eval)
   time_callback.on_train_end()
 
   stats = build_stats(runnable, time_callback)

official/vision/image_classification/resnet_imagenet_main.py → official/vision/image_classification/resnet/resnet_imagenet_main.py

@@ -26,16 +26,15 @@ from absl import logging
 import tensorflow as tf
 
 import tensorflow_model_optimization as tfmot
-
-from official.benchmark.models import trivial_model
 from official.modeling import performance
 from official.utils.flags import core as flags_core
 from official.utils.logs import logger
 from official.utils.misc import distribution_utils
 from official.utils.misc import keras_utils
 from official.utils.misc import model_helpers
-from official.vision.image_classification import common
-from official.vision.image_classification import imagenet_preprocessing
+from official.vision.image_classification import test_utils
+from official.vision.image_classification.resnet import common
+from official.vision.image_classification.resnet import imagenet_preprocessing
 from official.vision.image_classification.resnet import resnet_model
 
 
@@ -180,8 +179,7 @@ def run(flags_obj):
 
     # TODO(hongkuny): Remove trivial model usage and move it to benchmark.
     if flags_obj.use_trivial_model:
-      model = trivial_model.trivial_model(
-          imagenet_preprocessing.NUM_CLASSES)
+      model = test_utils.trivial_model(imagenet_preprocessing.NUM_CLASSES)
     elif flags_obj.model == 'resnet50_v1.5':
       model = resnet_model.resnet50(
           num_classes=imagenet_preprocessing.NUM_CLASSES)

official/vision/image_classification/resnet/resnet_model.py

@@ -33,7 +33,7 @@ from tensorflow.python.keras import backend
 from tensorflow.python.keras import initializers
 from tensorflow.python.keras import models
 from tensorflow.python.keras import regularizers
-from official.vision.image_classification import imagenet_preprocessing
+from official.vision.image_classification.resnet import imagenet_preprocessing
 
 L2_WEIGHT_DECAY = 1e-4
 BATCH_NORM_DECAY = 0.9
@@ -255,9 +255,7 @@ def resnet50(num_classes,
     x = img_input
 
   if backend.image_data_format() == 'channels_first':
-    x = layers.Lambda(
-        lambda x: backend.permute_dimensions(x, (0, 3, 1, 2)),
-        name='transpose')(x)
+    x = layers.Permute((3, 1, 2))(x)
     bn_axis = 1
   else:  # channels_last
     bn_axis = 3
@@ -382,8 +380,7 @@ def resnet50(num_classes,
       block='c',
       use_l2_regularizer=use_l2_regularizer)
 
-  rm_axes = [1, 2] if backend.image_data_format() == 'channels_last' else [2, 3]
-  x = layers.Lambda(lambda x: backend.mean(x, rm_axes), name='reduce_mean')(x)
+  x = layers.GlobalAveragePooling2D()(x)
   x = layers.Dense(
       num_classes,
       kernel_initializer=initializers.RandomNormal(stddev=0.01),

official/vision/image_classification/resnet_runnable.py → official/vision/image_classification/resnet/resnet_runnable.py

@@ -25,8 +25,8 @@ from official.staging.training import grad_utils
 from official.staging.training import standard_runnable
 from official.staging.training import utils
 from official.utils.flags import core as flags_core
-from official.vision.image_classification import common
-from official.vision.image_classification import imagenet_preprocessing
+from official.vision.image_classification.resnet import common
+from official.vision.image_classification.resnet import imagenet_preprocessing
 from official.vision.image_classification.resnet import resnet_model
 
 
@@ -175,7 +175,7 @@ class ResnetRunnable(standard_runnable.StandardTrainable,
       self.train_loss.update_state(loss)
       self.train_accuracy.update_state(labels, logits)
 
-    self.strategy.experimental_run_v2(step_fn, args=(next(iterator),))
+    self.strategy.run(step_fn, args=(next(iterator),))
 
   def train_loop_end(self):
     """See base class."""
@@ -204,7 +204,7 @@ class ResnetRunnable(standard_runnable.StandardTrainable,
       self.test_loss.update_state(loss)
       self.test_accuracy.update_state(labels, logits)
 
-    self.strategy.experimental_run_v2(step_fn, args=(next(iterator),))
+    self.strategy.run(step_fn, args=(next(iterator),))
 
   def eval_end(self):
     """See base class."""

official/vision/image_classification/resnet/tfhub_export.py

@@ -26,7 +26,7 @@ from absl import flags
 
 import tensorflow.compat.v2 as tf
 
-from official.vision.image_classification import imagenet_preprocessing
+from official.vision.image_classification.resnet import imagenet_preprocessing
 from official.vision.image_classification.resnet import resnet_model
 
 FLAGS = flags.FLAGS

official/vision/image_classification/resnet_ctl_imagenet_test.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.
-# ==============================================================================
-"""Test the ResNet model with ImageNet data using CTL."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import tempfile
-import os
-
-import tensorflow.compat.v2 as tf
-
-from tensorflow.python.eager import context
-from official.utils.testing import integration
-from official.vision.image_classification import common
-from official.vision.image_classification import imagenet_preprocessing
-from official.vision.image_classification import resnet_ctl_imagenet_main
-
-
-class CtlImagenetTest(tf.test.TestCase):
-  """Unit tests for Keras ResNet with ImageNet using CTL."""
-
-  _extra_flags = [
-      '-batch_size', '4',
-      '-train_steps', '4',
-      '-use_synthetic_data', 'true'
-  ]
-  _tempdir = None
-
-  def get_temp_dir(self):
-    if not self._tempdir:
-      self._tempdir = tempfile.mkdtemp(
-          dir=super(CtlImagenetTest, self).get_temp_dir())
-    return self._tempdir
-
-  @classmethod
-  def setUpClass(cls):
-    super(CtlImagenetTest, cls).setUpClass()
-    common.define_keras_flags()
-
-  def setUp(self):
-    super(CtlImagenetTest, self).setUp()
-    imagenet_preprocessing.NUM_IMAGES['validation'] = 4
-    self.policy = \
-        tf.compat.v2.keras.mixed_precision.experimental.global_policy()
-
-  def tearDown(self):
-    super(CtlImagenetTest, self).tearDown()
-    tf.io.gfile.rmtree(self.get_temp_dir())
-    tf.compat.v2.keras.mixed_precision.experimental.set_policy(self.policy)
-
-  def test_end_to_end_no_dist_strat(self):
-    """Test Keras model with 1 GPU, no distribution strategy."""
-
-    model_dir = os.path.join(self.get_temp_dir(), 'ctl_imagenet_no_dist_strat')
-    extra_flags = [
-        '-distribution_strategy', 'off',
-        '-model_dir', model_dir,
-        '-data_format', 'channels_last',
-    ]
-    extra_flags = extra_flags + self._extra_flags
-
-    integration.run_synthetic(
-        main=resnet_ctl_imagenet_main.run,
-        tmp_root=self.get_temp_dir(),
-        extra_flags=extra_flags
-    )
-
-  def test_end_to_end_2_gpu(self):
-    """Test Keras model with 2 GPUs."""
-    num_gpus = '2'
-    if context.num_gpus() < 2:
-      num_gpus = '0'
-
-    model_dir = os.path.join(self.get_temp_dir(), 'ctl_imagenet_2_gpu')
-    extra_flags = [
-        '-num_gpus', num_gpus,
-        '-distribution_strategy', 'mirrored',
-        '-model_dir', model_dir,
-        '-data_format', 'channels_last',
-    ]
-    extra_flags = extra_flags + self._extra_flags
-
-    integration.run_synthetic(
-        main=resnet_ctl_imagenet_main.run,
-        tmp_root=self.get_temp_dir(),
-        extra_flags=extra_flags
-    )
-
-if __name__ == '__main__':
-  assert tf.version.VERSION.startswith('2.')
-  tf.test.main()

official/vision/image_classification/resnet_imagenet_test.py

-# Copyright 2017 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.
-# ==============================================================================
-"""Test the keras ResNet model with ImageNet data."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from absl.testing import parameterized
-import tensorflow as tf
-
-from tensorflow.python.eager import context
-from official.utils.misc import keras_utils
-from official.utils.testing import integration
-from official.vision.image_classification import imagenet_preprocessing
-from official.vision.image_classification import resnet_imagenet_main
-
-
-@parameterized.parameters(
-    "resnet",
-    "resnet_polynomial_decay",
-    "mobilenet",
-    "mobilenet_polynomial_decay")
-class KerasImagenetTest(tf.test.TestCase):
-  """Unit tests for Keras Models with ImageNet."""
-  _default_flags_dict = [
-      "-batch_size", "4",
-      "-train_steps", "1",
-      "-use_synthetic_data", "true",
-      "-data_format", "channels_last",
-  ]
-  _extra_flags_dict = {
-      "resnet": [
-          "-model", "resnet50_v1.5",
-          "-optimizer", "resnet50_default",
-      ],
-      "resnet_polynomial_decay": [
-          "-model", "resnet50_v1.5",
-          "-optimizer", "resnet50_default",
-          "-pruning_method", "polynomial_decay",
-      ],
-      "mobilenet": [
-          "-model", "mobilenet",
-          "-optimizer", "mobilenet_default",
-      ],
-      "mobilenet_polynomial_decay": [
-          "-model", "mobilenet",
-          "-optimizer", "mobilenet_default",
-          "-pruning_method", "polynomial_decay",
-      ],
-  }
-  _tempdir = None
-
-  @classmethod
-  def setUpClass(cls):  # pylint: disable=invalid-name
-    super(KerasImagenetTest, cls).setUpClass()
-    resnet_imagenet_main.define_imagenet_keras_flags()
-
-  def setUp(self):
-    super(KerasImagenetTest, self).setUp()
-    imagenet_preprocessing.NUM_IMAGES["validation"] = 4
-    self.policy = \
-        tf.compat.v2.keras.mixed_precision.experimental.global_policy()
-
-  def tearDown(self):
-    super(KerasImagenetTest, self).tearDown()
-    tf.io.gfile.rmtree(self.get_temp_dir())
-    tf.compat.v2.keras.mixed_precision.experimental.set_policy(self.policy)
-
-  def get_extra_flags_dict(self, flags_key):
-    return self._extra_flags_dict[flags_key] + self._default_flags_dict
-
-  def test_end_to_end_no_dist_strat(self, flags_key):
-    """Test Keras model with 1 GPU, no distribution strategy."""
-    config = keras_utils.get_config_proto_v1()
-    tf.compat.v1.enable_eager_execution(config=config)
-
-    extra_flags = [
-        "-distribution_strategy", "off",
-    ]
-    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)
-
-    integration.run_synthetic(
-        main=resnet_imagenet_main.run,
-        tmp_root=self.get_temp_dir(),
-        extra_flags=extra_flags
-    )
-
-  def test_end_to_end_graph_no_dist_strat(self, flags_key):
-    """Test Keras model in legacy graph mode with 1 GPU, no dist strat."""
-    extra_flags = [
-        "-enable_eager", "false",
-        "-distribution_strategy", "off",
-    ]
-    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)
-
-    integration.run_synthetic(
-        main=resnet_imagenet_main.run,
-        tmp_root=self.get_temp_dir(),
-        extra_flags=extra_flags
-    )
-
-  def test_end_to_end_1_gpu(self, flags_key):
-    """Test Keras model with 1 GPU."""
-    config = keras_utils.get_config_proto_v1()
-    tf.compat.v1.enable_eager_execution(config=config)
-
-    if context.num_gpus() < 1:
-      self.skipTest(
-          "{} GPUs are not available for this test. {} GPUs are available".
-          format(1, context.num_gpus()))
-
-    extra_flags = [
-        "-num_gpus", "1",
-        "-distribution_strategy", "mirrored",
-        "-enable_checkpoint_and_export", "1",
-    ]
-    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)
-
-    integration.run_synthetic(
-        main=resnet_imagenet_main.run,
-        tmp_root=self.get_temp_dir(),
-        extra_flags=extra_flags
-    )
-
-  def test_end_to_end_1_gpu_fp16(self, flags_key):
-    """Test Keras model with 1 GPU and fp16."""
-    config = keras_utils.get_config_proto_v1()
-    tf.compat.v1.enable_eager_execution(config=config)
-
-    if context.num_gpus() < 1:
-      self.skipTest(
-          "{} GPUs are not available for this test. {} GPUs are available"
-          .format(1, context.num_gpus()))
-
-    extra_flags = [
-        "-num_gpus", "1",
-        "-dtype", "fp16",
-        "-distribution_strategy", "mirrored",
-    ]
-    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)
-
-    if "polynomial_decay" in extra_flags:
-      self.skipTest("Pruning with fp16 is not currently supported.")
-
-    integration.run_synthetic(
-        main=resnet_imagenet_main.run,
-        tmp_root=self.get_temp_dir(),
-        extra_flags=extra_flags
-    )
-
-  def test_end_to_end_2_gpu(self, flags_key):
-    """Test Keras model with 2 GPUs."""
-    config = keras_utils.get_config_proto_v1()
-    tf.compat.v1.enable_eager_execution(config=config)
-
-    if context.num_gpus() < 2:
-      self.skipTest(
-          "{} GPUs are not available for this test. {} GPUs are available".
-          format(2, context.num_gpus()))
-
-    extra_flags = [
-        "-num_gpus", "2",
-        "-distribution_strategy", "mirrored",
-    ]
-    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)
-
-    integration.run_synthetic(
-        main=resnet_imagenet_main.run,
-        tmp_root=self.get_temp_dir(),
-        extra_flags=extra_flags
-    )
-
-  def test_end_to_end_xla_2_gpu(self, flags_key):
-    """Test Keras model with XLA and 2 GPUs."""
-    config = keras_utils.get_config_proto_v1()
-    tf.compat.v1.enable_eager_execution(config=config)
-
-    if context.num_gpus() < 2:
-      self.skipTest(
-          "{} GPUs are not available for this test. {} GPUs are available".
-          format(2, context.num_gpus()))
-
-    extra_flags = [
-        "-num_gpus", "2",
-        "-enable_xla", "true",
-        "-distribution_strategy", "mirrored",
-    ]
-    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)
-
-    integration.run_synthetic(
-        main=resnet_imagenet_main.run,
-        tmp_root=self.get_temp_dir(),
-        extra_flags=extra_flags
-    )
-
-  def test_end_to_end_2_gpu_fp16(self, flags_key):
-    """Test Keras model with 2 GPUs and fp16."""
-    config = keras_utils.get_config_proto_v1()
-    tf.compat.v1.enable_eager_execution(config=config)
-
-    if context.num_gpus() < 2:
-      self.skipTest(
-          "{} GPUs are not available for this test. {} GPUs are available".
-          format(2, context.num_gpus()))
-
-    extra_flags = [
-        "-num_gpus", "2",
-        "-dtype", "fp16",
-        "-distribution_strategy", "mirrored",
-    ]
-    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)
-
-    if "polynomial_decay" in extra_flags:
-      self.skipTest("Pruning with fp16 is not currently supported.")
-
-    integration.run_synthetic(
-        main=resnet_imagenet_main.run,
-        tmp_root=self.get_temp_dir(),
-        extra_flags=extra_flags
-    )
-
-  def test_end_to_end_xla_2_gpu_fp16(self, flags_key):
-    """Test Keras model with XLA, 2 GPUs and fp16."""
-    config = keras_utils.get_config_proto_v1()
-    tf.compat.v1.enable_eager_execution(config=config)
-
-    if context.num_gpus() < 2:
-      self.skipTest(
-          "{} GPUs are not available for this test. {} GPUs are available".
-          format(2, context.num_gpus()))
-
-    extra_flags = [
-        "-num_gpus", "2",
-        "-dtype", "fp16",
-        "-enable_xla", "true",
-        "-distribution_strategy", "mirrored",
-    ]
-    extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)
-
-    if "polynomial_decay" in extra_flags:
-      self.skipTest("Pruning with fp16 is not currently supported.")
-
-    integration.run_synthetic(
-        main=resnet_imagenet_main.run,
-        tmp_root=self.get_temp_dir(),
-        extra_flags=extra_flags
-    )
-
-
-if __name__ == "__main__":
-  tf.compat.v1.enable_v2_behavior()
-  tf.test.main()

official/benchmark/models/trivial_model.py → official/vision/image_classification/test_utils.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""A trivial model for Keras."""
+"""Test utilities for image classification tasks."""
 
 from __future__ import absolute_import
 from __future__ import division

tutorials/README.md

-# Tutorial Models
-
-This folder contains models referenced to from the [TensorFlow tutorials](https://www.tensorflow.org/tutorials/).

tutorials/embedding/README.md

-This directory contains models for unsupervised training of word embeddings
-using the model described in:
-
-(Mikolov, et. al.) [Efficient Estimation of Word Representations in Vector Space](http://arxiv.org/abs/1301.3781),
-ICLR 2013.
-
-Detailed instructions on how to get started and use them are available in the
-tutorials. Brief instructions are below.
-
-* [Word2Vec Tutorial](http://tensorflow.org/tutorials/word2vec)
-
-Assuming you have cloned the git repository, navigate into this directory. To download the example text and evaluation data:
-
-```shell
-curl http://mattmahoney.net/dc/text8.zip > text8.zip
-unzip text8.zip
-curl https://storage.googleapis.com/google-code-archive-source/v2/code.google.com/word2vec/source-archive.zip > source-archive.zip
-unzip -p source-archive.zip  word2vec/trunk/questions-words.txt > questions-words.txt
-rm text8.zip source-archive.zip
-```
-
-You will need to compile the ops as follows (See 
-[Adding a New Op to TensorFlow](https://www.tensorflow.org/how_tos/adding_an_op/#building_the_op_library)
-for more details).:
-
-```shell
-TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') )
-TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))') )
-g++ -std=c++11 -shared word2vec_ops.cc word2vec_kernels.cc -o word2vec_ops.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2 -D_GLIBCXX_USE_CXX11_ABI=0
-```
-
-On Mac, add `-undefined dynamic_lookup` to the g++ command. The flag `-D_GLIBCXX_USE_CXX11_ABI=0` is included to support newer versions of gcc. However, if you compiled TensorFlow from source using gcc 5 or later, you may need to exclude the flag. Specifically, if you get an error similar to the following: `word2vec_ops.so: undefined symbol: _ZN10tensorflow7strings6StrCatERKNS0_8AlphaNumES3_S3_S3_` then you likely need to exclude the flag.
-
-Once you've successfully compiled the ops, run the model as follows:
-
-```shell
-python word2vec_optimized.py \
-  --train_data=text8 \
-  --eval_data=questions-words.txt \
-  --save_path=/tmp/
-```
-
-Here is a short overview of what is in this directory.
-
-File | What's in it?
---- | ---
-`word2vec.py` | A version of word2vec implemented using TensorFlow ops and minibatching.
-`word2vec_test.py` | Integration test for word2vec.
-`word2vec_optimized.py` | A version of word2vec implemented using C ops that does no minibatching.
-`word2vec_optimized_test.py` | Integration test for word2vec_optimized.
-`word2vec_kernels.cc` | Kernels for the custom input and training ops.
-`word2vec_ops.cc` | The declarations of the custom ops.

tutorials/embedding/__init__.py

-# Copyright 2015 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.
-# ==============================================================================
-
-"""Import generated word2vec optimized ops into embedding package."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function