Use nonexperimental mixed precision API.

This replaces symbols in tf.keras.mixed_precision.experimental with the corresponding nonexperimental symbols. In some cases, passing a Policy is replaced with passing a policy name for conciseness.

Additionally, for the Shakespeare model, the loss_scale flag is removed, since supporting it with the nonexperimental API is slightly more verbose and it is recommended users use the default loss scale.

PiperOrigin-RevId: 368123944

research/object_detection/model_lib.py

@@ -445,8 +445,7 @@ def create_model_fn(detection_model_fn, configs, hparams=None, use_tpu=False,
       from tensorflow.python.keras.engine import base_layer_utils  # pylint: disable=g-import-not-at-top
       # Enable v2 behavior, as `mixed_bfloat16` is only supported in TF 2.0.
       base_layer_utils.enable_v2_dtype_behavior()
-      tf2.keras.mixed_precision.experimental.set_policy(
-          'mixed_bfloat16')
+      tf2.keras.mixed_precision.set_global_policy('mixed_bfloat16')
     detection_model = detection_model_fn(
         is_training=is_training, add_summaries=(not use_tpu))
     scaffold_fn = None

research/object_detection/model_lib_v2.py

@@ -507,7 +507,7 @@ def train_loop(
     train_steps = train_config.num_steps
 
   if kwargs['use_bfloat16']:
-    tf.compat.v2.keras.mixed_precision.experimental.set_policy('mixed_bfloat16')
+    tf.compat.v2.keras.mixed_precision.set_global_policy('mixed_bfloat16')
 
   if train_config.load_all_detection_checkpoint_vars:
     raise ValueError('train_pb2.load_all_detection_checkpoint_vars '
@@ -1085,7 +1085,7 @@ def eval_continuously(
     eval_on_train_input_config.num_epochs = 1
 
   if kwargs['use_bfloat16']:
-    tf.compat.v2.keras.mixed_precision.experimental.set_policy('mixed_bfloat16')
+    tf.compat.v2.keras.mixed_precision.set_global_policy('mixed_bfloat16')
 
   eval_input_config = eval_input_configs[eval_index]
   strategy = tf.compat.v2.distribute.get_strategy()