Use new mixed_float16 policy for resnet.

The old infer_float32_policies policy will be removed from TensorFlow soon.

To test convergence, I ran the Resnet50KerasAccuracy.benchmark_8_gpu_fp16 benchmark. I got an accuracy of 0.76037 and an exp_per_second of 6908.

PiperOrigin-RevId: 266191126

official/benchmark/models/trivial_model.py

@@ -23,19 +23,16 @@ from tensorflow.python.keras import layers
 from tensorflow.python.keras import models
 
 
-def trivial_model(num_classes, dtype='float32'):
+def trivial_model(num_classes):
   """Trivial model for ImageNet dataset."""
 
   input_shape = (224, 224, 3)
-  img_input = layers.Input(shape=input_shape, dtype=dtype)
+  img_input = layers.Input(shape=input_shape)
 
   x = layers.Lambda(lambda x: backend.reshape(x, [-1, 224 * 224 * 3]),
                     name='reshape')(img_input)
   x = layers.Dense(1, name='fc1')(x)
   x = layers.Dense(num_classes, name='fc1000')(x)
-  # TODO(reedwm): Remove manual casts once mixed precision can be enabled with a
-  # single line of code.
-  x = backend.cast(x, 'float32')
-  x = layers.Activation('softmax')(x)
+  x = layers.Activation('softmax', dtype='float32')(x)
 
   return models.Model(img_input, x, name='trivial')

official/resnet/ctl/ctl_imagenet_main.py

@@ -141,8 +141,6 @@ def run(flags_obj):
       enable_eager=flags_obj.enable_eager,
       enable_xla=flags_obj.enable_xla)
 
-  dtype = flags_core.get_tf_dtype(flags_obj)
-
   # TODO(anj-s): Set data_format without using Keras.
   data_format = flags_obj.data_format
   if data_format is None:
@@ -167,7 +165,7 @@ def run(flags_obj):
   with strategy_scope:
     model = resnet_model.resnet50(
         num_classes=imagenet_preprocessing.NUM_CLASSES,
-        dtype=dtype, batch_size=flags_obj.batch_size,
+        batch_size=flags_obj.batch_size,
         use_l2_regularizer=not flags_obj.single_l2_loss_op)
 
     optimizer = tf.keras.optimizers.SGD(

official/vision/image_classification/resnet_imagenet_main.py

@@ -95,8 +95,12 @@ def run(flags_obj):
 
   dtype = flags_core.get_tf_dtype(flags_obj)
   if dtype == 'float16':
-    policy = tf.keras.mixed_precision.experimental.Policy('infer_float32_vars')
+    loss_scale = flags_core.get_loss_scale(flags_obj, default_for_fp16=128)
+    policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16',
+                                                          loss_scale=loss_scale)
     tf.keras.mixed_precision.experimental.set_policy(policy)
+    if not keras_utils.is_v2_0():
+      raise ValueError('--dtype=fp16 is not supported in TensorFlow 1.')
 
   data_format = flags_obj.data_format
   if data_format is None:
@@ -178,13 +182,6 @@ def run(flags_obj):
 
   with strategy_scope:
     optimizer = common.get_optimizer(lr_schedule)
-    if dtype == 'float16':
-      # TODO(reedwm): Remove manually wrapping optimizer once mixed precision
-      # can be enabled with a single line of code.
-      optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(
-          optimizer, loss_scale=flags_core.get_loss_scale(flags_obj,
-                                                          default_for_fp16=128))
-
     if flags_obj.fp16_implementation == "graph_rewrite":
       # Note: when flags_obj.fp16_implementation == "graph_rewrite", 
       # dtype as determined by flags_core.get_tf_dtype(flags_obj) would be 'float32'
@@ -195,10 +192,10 @@ 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, dtype)
+          imagenet_preprocessing.NUM_CLASSES)
     else:
       model = resnet_model.resnet50(
-          num_classes=imagenet_preprocessing.NUM_CLASSES, dtype=dtype)
+          num_classes=imagenet_preprocessing.NUM_CLASSES)
 
     # TODO(b/138957587): Remove when force_v2_in_keras_compile is on longer
     # a valid arg for this model. Also remove as a valid flag.

official/vision/image_classification/resnet_imagenet_test.py

@@ -116,6 +116,31 @@ class KerasImagenetTest(googletest.TestCase):
         extra_flags=extra_flags
     )
 
+  def test_end_to_end_1_gpu_fp16(self):
+    """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", "default",
+        "-model_dir", "keras_imagenet_1_gpu",
+        "-data_format", "channels_last",
+    ]
+    extra_flags = extra_flags + self._extra_flags
+
+    integration.run_synthetic(
+        main=resnet_imagenet_main.run,
+        tmp_root=self.get_temp_dir(),
+        extra_flags=extra_flags
+    )
+
   def test_end_to_end_graph_1_gpu(self):
     """Test Keras model in legacy graph mode with 1 GPU."""
     if context.num_gpus() < 1:
@@ -279,4 +304,5 @@ class KerasImagenetTest(googletest.TestCase):
 
 
 if __name__ == "__main__":
+  tf.compat.v1.enable_v2_behavior()
   googletest.main()

official/vision/image_classification/resnet_model.py

@@ -222,14 +222,12 @@ def conv_block(input_tensor,
 
 
 def resnet50(num_classes,
-             dtype='float32',
              batch_size=None,
              use_l2_regularizer=True):
   """Instantiates the ResNet50 architecture.
 
   Args:
     num_classes: `int` number of classes for image classification.
-    dtype: dtype to use float32 or float16 are most common.
     batch_size: Size of the batches for each step.
     use_l2_regularizer: whether to use L2 regularizer on Conv/Dense layer.
 
@@ -237,8 +235,7 @@ def resnet50(num_classes,
       A Keras model instance.
   """
   input_shape = (224, 224, 3)
-  img_input = layers.Input(
-      shape=input_shape, dtype=dtype, batch_size=batch_size)
+  img_input = layers.Input(shape=input_shape, batch_size=batch_size)
 
   if backend.image_data_format() == 'channels_first':
     x = layers.Lambda(
@@ -380,10 +377,9 @@ def resnet50(num_classes,
       name='fc1000')(
           x)
 
-  # TODO(reedwm): Remove manual casts once mixed precision can be enabled with a
-  # single line of code.
-  x = backend.cast(x, 'float32')
-  x = layers.Activation('softmax')(x)
+  # A softmax that is followed by the model loss must be done cannot be done
+  # in float16 due to numeric issues. So we pass dtype=float32.
+  x = layers.Activation('softmax', dtype='float32')(x)
 
   # Create model.
   return models.Model(img_input, x, name='resnet50')