def __init__(self,
hidden_1_dim=256,
hidden_2_dim=256,
input_dim=28 * 28,
classes=10):
super().__init__()
tf_test_utils.set_random_seed()
self.hidden_1_dim = hidden_1_dim
self.hidden_2_dim = hidden_2_dim
self.input_dim = input_dim
self.classes = classes
self.h1_weights = tf.Variable(
tf.random.normal([input_dim, hidden_1_dim]))
self.h2_weights = tf.Variable(
tf.random.normal([hidden_1_dim, hidden_2_dim]))
self.out_weights = tf.Variable(
tf.random.normal([hidden_2_dim, classes]))
self.h1_bias = tf.Variable(tf.random.normal([hidden_1_dim]))
self.h2_bias = tf.Variable(tf.random.normal([hidden_2_dim]))
self.out_bias = tf.Variable(tf.random.normal([classes]))
# Compile with dynamic batch dim.
self.predict = tf.function(
input_signature=[tf.TensorSpec([None, self.input_dim])])(
self.predict)
def models():
tf.keras.backend.set_learning_phase(False)
tf_test_utils.set_random_seed()
input_shape = get_input_shape(FLAGS.data)
# keras model receives images size as input,
# where batch size is not specified - by default it is dynamic
if FLAGS.model in APP_MODELS:
weights = 'imagenet' if FLAGS.data == 'imagenet' else None
# TODO(rybakov) with the fix of https://github.com/google/iree/issues/1660
# add include_top=True
# if weights == 'imagenet' it will load weights from external tf.keras URL
model = APP_MODELS[FLAGS.model](weights=weights,
input_shape=input_shape[1:])
if FLAGS.data == 'cifar10' and FLAGS.url:
file_name = 'cifar10' + FLAGS.model
# it will download model weights from publically available folder: PATH
# and save it to cache_dir=~/.keras and return path to it
weights_path = tf.keras.utils.get_file(
file_name, os.path.join(FLAGS.url, file_name + '.h5'))
model.load_weights(weights_path)
else:
raise ValueError('Unsupported model', FLAGS.model)
module = tf.Module()
module.m = model
# specify input size with static batch size
# TODO(b/142948097): with support of dynamic shape
# replace input_shape by model.input_shape, so batch size will be dynamic (-1)
module.predict = tf.function(input_signature=[tf.TensorSpec(input_shape)])(
model.call)
return module
def lstm_module():
tf_test_utils.set_random_seed()
inputs = tf.keras.layers.Input(batch_size=NUM_BATCH, shape=INPUT_SHAPE[1:])
outputs = tf.keras.layers.LSTM(units=NUM_UNITS,
return_sequences=True)(inputs)
model = tf.keras.Model(inputs, outputs)
module = tf.Module()
module.m = model
module.predict = tf.function(
input_signature=[tf.TensorSpec(INPUT_SHAPE, tf.float32)])(model.call)
return module
def CreateModule(input_dim=_FEATURE_SIZE, output_dim=1):
"""Creates a module for regression model training.
Args:
input_dim: input dimensionality
output_dim: output dimensionality
Returns:
model for linear regression
"""
tf_test_utils.set_random_seed()
# build a single layer model
inputs = tf.keras.layers.Input((input_dim))
outputs = tf.keras.layers.Dense(output_dim)(inputs)
model = tf.keras.Model(inputs, outputs)
return ModelTrain(model)
def models():
tf.keras.backend.set_learning_phase(False)
# TODO(ataei): This should move somewhere in SavedModelTestCase, it should
# guarantee test is deterministic.
tf_test_utils.set_random_seed()
# keras model receives images size as input,
# where batch size is not specified - by default it is dynamic
if FLAGS.model in APP_MODELS:
model = APP_MODELS[FLAGS.model](
weights=None, include_top=False, input_shape=INPUT_SHAPE[1:])
else:
raise ValueError('unsupported model', FLAGS.model)
module = tf.Module()
module.m = model
# specify input size with static batch size
# TODO(b/142948097): with support of dynamic shape
# replace INPUT_SHAPE by model.input_shape, so batch size will be dynamic (-1)
module.predict = tf.function(input_signature=[tf.TensorSpec(INPUT_SHAPE)])(
model.call)
return module
