Internal change

PiperOrigin-RevId: 264895439

official/resnet/ctl/ctl_imagenet_benchmark.py

@@ -22,7 +22,7 @@ import time
 from absl import flags
 import tensorflow as tf
 
-from official.resnet.keras import keras_common
+from official.vision.image_classification import common
 from official.resnet.ctl import ctl_imagenet_main
 from official.resnet.ctl import ctl_common
 from official.utils.testing.perfzero_benchmark import PerfZeroBenchmark
@@ -118,7 +118,7 @@ class Resnet50CtlAccuracy(CtlBenchmark):
 
     flag_methods = [
         ctl_common.define_ctl_flags,
-        keras_common.define_keras_flags
+        common.define_keras_flags
     ]
 
     self.data_dir = os.path.join(root_data_dir, 'imagenet')
@@ -162,7 +162,7 @@ class Resnet50CtlBenchmarkBase(CtlBenchmark):
   def __init__(self, output_dir=None, default_flags=None):
     flag_methods = [
         ctl_common.define_ctl_flags,
-        keras_common.define_keras_flags
+        common.define_keras_flags
     ]
 
     super(Resnet50CtlBenchmarkBase, self).__init__(

official/resnet/ctl/ctl_imagenet_main.py

@@ -24,10 +24,10 @@ from absl import logging
 import tensorflow as tf
 
 from official.resnet.ctl import ctl_common
-from official.resnet.keras import imagenet_preprocessing
-from official.resnet.keras import keras_common
-from official.resnet.keras import keras_imagenet_main
-from official.resnet.keras import resnet_model
+from official.vision.image_classification import imagenet_preprocessing
+from official.vision.image_classification import common
+from official.vision.image_classification import resnet_imagenet_main
+from official.vision.image_classification import resnet_model
 from official.utils.flags import core as flags_core
 from official.utils.logs import logger
 from official.utils.misc import distribution_utils
@@ -73,7 +73,7 @@ def get_input_dataset(flags_obj, strategy):
   """Returns the test and train input datasets."""
   dtype = flags_core.get_tf_dtype(flags_obj)
   if flags_obj.use_synthetic_data:
-    input_fn = keras_common.get_synth_input_fn(
+    input_fn = common.get_synth_input_fn(
         height=imagenet_preprocessing.DEFAULT_IMAGE_SIZE,
         width=imagenet_preprocessing.DEFAULT_IMAGE_SIZE,
         num_channels=imagenet_preprocessing.NUM_CHANNELS,
@@ -171,7 +171,7 @@ def run(flags_obj):
         use_l2_regularizer=not flags_obj.single_l2_loss_op)
 
     optimizer = tf.keras.optimizers.SGD(
-        learning_rate=keras_common.BASE_LEARNING_RATE, momentum=0.9,
+        learning_rate=common.BASE_LEARNING_RATE, momentum=0.9,
         nesterov=True)
 
     training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
@@ -248,7 +248,7 @@ def run(flags_obj):
       training_accuracy.reset_states()
 
       for step in range(train_steps):
-        optimizer.lr = keras_imagenet_main.learning_rate_schedule(
+        optimizer.lr = resnet_imagenet_main.learning_rate_schedule(
             epoch, step, train_steps, flags_obj.batch_size)
 
         time_callback.on_batch_begin(step+epoch*train_steps)
@@ -297,7 +297,7 @@ def main(_):
 
 if __name__ == '__main__':
   logging.set_verbosity(logging.INFO)
-  keras_common.define_keras_flags()
+  common.define_keras_flags()
   ctl_common.define_ctl_flags()
   flags.adopt_module_key_flags(keras_common)
   flags.adopt_module_key_flags(ctl_common)

official/resnet/ctl/ctl_imagenet_test.py

@@ -25,8 +25,8 @@ from tensorflow.python.eager import context
 from tensorflow.python.platform import googletest
 from official.resnet.ctl import ctl_common
 from official.resnet.ctl import ctl_imagenet_main
-from official.resnet.keras import imagenet_preprocessing
-from official.resnet.keras import keras_common
+from official.vision.image_classification import imagenet_preprocessing
+from official.vision.image_classification import common
 from official.utils.misc import keras_utils
 from official.utils.testing import integration
 
@@ -49,7 +49,7 @@ class CtlImagenetTest(googletest.TestCase):
   @classmethod
   def setUpClass(cls):  # pylint: disable=invalid-name
     super(CtlImagenetTest, cls).setUpClass()
-    keras_common.define_keras_flags()
+    common.define_keras_flags()
     ctl_common.define_ctl_flags()
 
   def setUp(self):

official/resnet/keras/__init__.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.
-# ==============================================================================
-
-"""Bring in the shared Keras ResNet modules into this module.
-
-The TensorFlow official Keras models are moved under
-    official/vision/image_classification
-In order to be backward compatible with models that directly import its modules,
-we import the Keras ResNet modules under official.resnet.keras.
-
-New TF models should not depend on modules directly under this path.
-"""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from official.vision.image_classification import cifar_preprocessing
-from official.vision.image_classification import common as keras_common
-from official.vision.image_classification import imagenet_preprocessing
-from official.vision.image_classification import resnet_cifar_main as keras_cifar_main
-from official.vision.image_classification import resnet_cifar_model
-from official.vision.image_classification import resnet_imagenet_main as keras_imagenet_main
-from official.vision.image_classification import resnet_model
-
-del absolute_import
-del division
-del print_function

official/vision/image_classification/resnet_imagenet_main.py

@@ -31,7 +31,7 @@ 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_model
-from official.vision.image_classification import trivial_model
+from official.benchmark.models import trivial_model
 
 
 LR_SCHEDULE = [    # (multiplier, epoch to start) tuples
@@ -186,6 +186,7 @@ def run(flags_obj):
           optimizer, loss_scale=flags_core.get_loss_scale(flags_obj,
                                                           default_for_fp16=128))
 
+    # 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, dtype)