Convert model_garden/vision code to use public tf.keras API.

PiperOrigin-RevId: 369893606

official/vision/image_classification/test_utils.py

@@ -18,21 +18,20 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow.python.keras import backend
-from tensorflow.python.keras import layers
-from tensorflow.python.keras import models
+import tensorflow as tf
 
 
 def trivial_model(num_classes):
   """Trivial model for ImageNet dataset."""
 
   input_shape = (224, 224, 3)
-  img_input = layers.Input(shape=input_shape)
+  img_input = tf.keras.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)
-  x = layers.Activation('softmax', dtype='float32')(x)
+  x = tf.keras.layers.Lambda(
+      lambda x: tf.keras.backend.reshape(x, [-1, 224 * 224 * 3]),
+      name='reshape')(img_input)
+  x = tf.keras.layers.Dense(1, name='fc1')(x)
+  x = tf.keras.layers.Dense(num_classes, name='fc1000')(x)
+  x = tf.keras.layers.Activation('softmax', dtype='float32')(x)
 
-  return models.Model(img_input, x, name='trivial')
+  return tf.keras.models.Model(img_input, x, name='trivial')