"vscode:/vscode.git/clone" did not exist on "4bcbe8e140a6c83c151b93567f33e60e486922bd"
Commit d54edbcf authored by Scott Zhu's avatar Scott Zhu Committed by A. Unique TensorFlower
Browse files

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

PiperOrigin-RevId: 369893606
parent 991f75e2
...@@ -18,21 +18,20 @@ from __future__ import absolute_import ...@@ -18,21 +18,20 @@ from __future__ import absolute_import
from __future__ import division from __future__ import division
from __future__ import print_function from __future__ import print_function
from tensorflow.python.keras import backend import tensorflow as tf
from tensorflow.python.keras import layers
from tensorflow.python.keras import models
def trivial_model(num_classes): def trivial_model(num_classes):
"""Trivial model for ImageNet dataset.""" """Trivial model for ImageNet dataset."""
input_shape = (224, 224, 3) 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]), x = tf.keras.layers.Lambda(
lambda x: tf.keras.backend.reshape(x, [-1, 224 * 224 * 3]),
name='reshape')(img_input) name='reshape')(img_input)
x = layers.Dense(1, name='fc1')(x) x = tf.keras.layers.Dense(1, name='fc1')(x)
x = layers.Dense(num_classes, name='fc1000')(x) x = tf.keras.layers.Dense(num_classes, name='fc1000')(x)
x = layers.Activation('softmax', dtype='float32')(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')
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment