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