def test_encode_decode_type_spec(self): spec = tf.TensorSpec((1, 5), tf.float32) string = json_utils.Encoder().encode(spec) loaded = json_utils.decode(string) self.assertEqual(spec, loaded) invalid_type_spec = { 'class_name': 'TypeSpec', 'type_spec': 'Invalid Type', 'serialized': None } string = json_utils.Encoder().encode(invalid_type_spec) with self.assertRaisesRegexp(ValueError, 'No TypeSpec has been registered'): loaded = json_utils.decode(string)
def testTypeSpecArg(self): # Create a Keras Input x = input_layer_lib.Input(type_spec=tf.TensorSpec((7, 32), tf.float32)) self.assertAllEqual(x.shape.as_list(), [7, 32]) # Verify you can construct and use a model w/ this input model = functional.Functional(x, x * 2.0) self.assertAllEqual(model(tf.ones(x.shape)), tf.ones(x.shape) * 2.0) # Test serialization / deserialization model = functional.Functional.from_config(model.get_config()) self.assertAllEqual(model(tf.ones(x.shape)), tf.ones(x.shape) * 2.0) model = model_config.model_from_json(model.to_json()) self.assertAllEqual(model(tf.ones(x.shape)), tf.ones(x.shape) * 2.0)
def create_serving_signature(model): @tf.function def serve_fn(raw_features): raw_features = tf.expand_dims(raw_features, axis=0) transformed_features = model.feature_ps(raw_features) outputs = model(transformed_features) outputs = tf.squeeze(outputs, axis=0) outputs = tf.cast(tf.greater(outputs, 0.5), tf.int64) decoded_outputs = model.label_inverse_lookup_layer(outputs) return tf.squeeze(decoded_outputs, axis=0) # serving does NOT have batch dimension return serve_fn.get_concrete_function( tf.TensorSpec(shape=(3), dtype=tf.string, name="example"))
class LanguageModelEncoder(tf.train.Checkpoint): def __init__(self,vocab_size,emb_dim,state_size,n_layers): super(LanguageModelEncoder, self).__init__() self._state_size = state_size self.embedding_layer = Embedding(vocab_size,emb_dim) self._lstm_layers = [LSTM(self._state_size,return_sequences=True) for i in range(n_layers)] @tf.function(input_signature=[tf.TensorSpec([None,None], tf.dtypes.int64)]) def __call__(self,sentence_lookup_ids): emb_output = self.embedding_layer(sentence_lookup_ids) lstm_output = emb_output # initialize to the input for lstm_layer in self._lstm_layers: lstm_output = lstm_layer(lstm_output) return lstm_output
class ControlFlowModule(tf.Module): def __init__(self): pass @tf.function(input_signature=[tf.TensorSpec([], tf.float32)]) def collatz(self, a): i = 0. while a > 1.: i = i + 1. if (a % 2.) > 0.: a = 3. * a + 1. else: a = a / 2. return i
def setUp(self): super().setUp() input_signature = [ tf.TensorSpec(shape=[None], dtype=tf.float32, name='init') ] self.init = tf.ones([7]) self.decorator = tf.function(jit_compile=True, input_signature=input_signature) self.arms = 2 self.days = 3 self.n_chains = 5 self.trials = tfd.Poisson(100.).sample([self.arms, self.days]) self.dist = get_joint_distribution(self.trials) self.true_values = self.dist.sample(seed=test_util.test_seed())
class TestModule(tf.Module): # Check that we get shapes annotated on function arguments. # # Besides checking the shape on the function input argument, this test also # checks that the shape on the input argument is propagated to the return # value. # We eventually want to move the shape inference to a pass separate from # the initial import, in which case that aspect of this test doesn't make much # sense and will be superceded by MLIR->MLIR shape inference tests. # # CHECK: func {{@[a-zA-Z_0-9]+}}(%arg0: tensor<f32> {{.*}}) -> (tensor<f32> {{.*}}) # CHECK-NEXT: attributes {{.*}} tf_saved_model.exported_names = ["some_function"] @tf.function(input_signature=[tf.TensorSpec([], tf.float32)]) def some_function(self, x): return x
def test_encode_decode_type_spec(self): spec = tf.TensorSpec((1, 5), tf.float32) string = json_utils.Encoder().encode(spec) loaded = json_utils.decode(string) self.assertEqual(spec, loaded) invalid_type_spec = { "class_name": "TypeSpec", "type_spec": "Invalid Type", "serialized": None, } string = json_utils.Encoder().encode(invalid_type_spec) with self.assertRaisesRegexp( ValueError, "No TypeSpec has been registered" ): loaded = json_utils.decode(string)
def test_specify_input_signature(self): model = testing_utils.get_small_sequential_mlp(10, 3, None) inputs = tf.ones((8, 5)) with self.assertRaisesRegex(ValueError, 'input shapes have not been set'): saving_utils.trace_model_call(model) fn = saving_utils.trace_model_call( model, [tf.TensorSpec(shape=[None, 5], dtype=tf.float32)]) signature_outputs = fn(inputs) if model.output_names: expected_outputs = {model.output_names[0]: model(inputs)} else: expected_outputs = {'output_1': model(inputs)} self._assert_all_close(expected_outputs, signature_outputs)
class TextEmbeddingModel(tf.train.Checkpoint): """Text embedding model. A text embeddings model that takes a sentences on input and outputs the sentence embedding. """ def __init__(self, vocabulary, emb_dim, oov_buckets): super(TextEmbeddingModel, self).__init__() self._oov_buckets = oov_buckets self._vocabulary_file = tracking.TrackableAsset( write_vocabulary_file(vocabulary)) self._total_size = len(vocabulary) + oov_buckets self._table = lookup_ops.index_table_from_file( vocabulary_file=self._vocabulary_file, num_oov_buckets=self._oov_buckets, hasher_spec=lookup_ops.FastHashSpec) self.embeddings = tf.Variable( tf.random.uniform(shape=[self._total_size, emb_dim])) self.variables = [self.embeddings] self.trainable_variables = self.variables def _tokenize(self, sentences): # Perform a minimalistic text preprocessing by removing punctuation and # splitting on spaces. normalized_sentences = tf.strings.regex_replace( input=sentences, pattern=r"\pP", rewrite="") normalized_sentences = tf.reshape(normalized_sentences, [-1]) sparse_tokens = tf.string_split(normalized_sentences, " ") # Deal with a corner case: there is one empty sentence. sparse_tokens, _ = tf.sparse.fill_empty_rows(sparse_tokens, tf.constant("")) # Deal with a corner case: all sentences are empty. sparse_tokens = tf.sparse.reset_shape(sparse_tokens) sparse_token_ids = self._table.lookup(sparse_tokens.values) return (sparse_tokens.indices, sparse_token_ids, sparse_tokens.dense_shape) @tf.function(input_signature=[tf.TensorSpec([None], tf.dtypes.string)]) def __call__(self, sentences): token_ids, token_values, token_dense_shape = self._tokenize(sentences) return tf.nn.safe_embedding_lookup_sparse( embedding_weights=self.embeddings, sparse_ids=tf.SparseTensor(token_ids, token_values, token_dense_shape), sparse_weights=None, combiner="sqrtn")
def test_save_load_io_device(self, model_and_input, distribution): saved_dir = os.path.join(self.get_temp_dir(), 'io_device') with distribution.scope(): model = model_and_input.get_model() x_train, y_train, _ = model_and_input.get_data() batch_size = model_and_input.get_batch_size() self._train_model(model, x_train, y_train, batch_size) call = model.__call__.get_concrete_function(tf.TensorSpec(None)) save_options = tf.saved_model.SaveOptions( experimental_io_device='/job:localhost') tf.saved_model.save(model, saved_dir, signatures=call, options=save_options) load_options = tf.saved_model.LoadOptions( experimental_io_device='/job:localhost') # Check that the model can be loaded and training continued without error. with distribution.scope(): loaded_model = tf.saved_model.load(saved_dir, options=load_options) self._train_model(loaded_model, x_train, y_train, batch_size)
class TestModule(tf.Module): def __init__(self): super(TestModule, self).__init__() self.v42 = tf.Variable(42.0) self.c43 = tf.constant(43.0) # CHECK: "tf_saved_model.global_tensor"() {is_mutable, sym_name = "[[VAR:[a-zA-Z_0-9]+]]", tf_saved_model.exported_names = ["v42"], type = tensor<f32>, value = dense<4.200000e+01> : tensor<f32>} : () -> () # CHECK: "tf_saved_model.global_tensor"() {sym_name = "[[CONST:[a-zA-Z_0-9]+]]", tf_saved_model.exported_names = [], type = tensor<f32>, value = dense<4.300000e+01> : tensor<f32>} : () -> () # CHECK: func {{@[a-zA-Z_0-9]+}}( # CHECK-SAME: %arg0: tensor<f32> {tf._user_specified_name = "x", tf_saved_model.index_path = [0]}, # CHECK-SAME: %arg1: tensor<!tf_type.resource<tensor<f32>>> {tf_saved_model.bound_input = @[[VAR]]}, # CHECK-SAME: %arg2: tensor<!tf_type.resource<tensor<f32>>> {tf_saved_model.bound_input = @[[CONST]]}) -> ( # CHECK-SAME: tensor<f32> {tf_saved_model.index_path = []}) # CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["some_function"] @tf.function(input_signature=[tf.TensorSpec([], tf.float32)]) def some_function(self, x): return x + self.v42 + self.c43
def testSaveSequentialModelWithoutInputShapes(self): model = sequential_model_without_input_shape(True) # A Sequential model that hasn't been built should raise an error. with self.assertRaisesRegex( ValueError, 'Weights for sequential model have not yet been created'): keras_saved_model.export_saved_model(model, '') # Even with input_signature, the model's weights has not been created. with self.assertRaisesRegex( ValueError, 'Weights for sequential model have not yet been created'): saved_model_dir = self._save_model_dir() keras_saved_model.export_saved_model( model, saved_model_dir, input_signature=tf.TensorSpec( shape=(10, 11, 12, 13, 14), dtype=tf.float32, name='spec_input'))
class TestGetTensorSpec(parameterized.TestCase): @parameterized.parameters([ (lambda: tf.constant([[1, 2]]), [1, 2]), (tf.TensorSpec([8, 3], tf.int32), [8, 3]), (tf.TensorSpec([8], tf.int32), [8]), (tf.TensorSpec([], tf.int32), []), (tf.TensorSpec(None, tf.int32), None), (tf.RaggedTensorSpec([8, 3], tf.int32), [8, 3]), (tf.SparseTensorSpec([8, 3], tf.int32), [8, 3]), ]) def test_without_dynamic_batch(self, t, expected_shape): if callable(t): t = t() result = tf_utils.get_tensor_spec(t) self.assertTrue(result.is_compatible_with(t)) if expected_shape is None: self.assertIsNone(result.shape.rank) else: self.assertEqual(result.shape.as_list(), expected_shape) @parameterized.parameters([ (lambda: tf.constant([[1, 2]]), [None, 2]), (tf.TensorSpec([8, 3], tf.int32), [None, 3]), (tf.TensorSpec([8], tf.int32), [None]), (tf.TensorSpec([], tf.int32), []), (tf.TensorSpec(None, tf.int32), None), (tf.RaggedTensorSpec([8, 3], tf.int32), [None, 3]), (tf.SparseTensorSpec([8, 3], tf.int32), [None, 3]), ]) def test_with_dynamic_batch(self, t, expected_shape): if callable(t): t = t() result = tf_utils.get_tensor_spec(t, True) self.assertTrue(result.is_compatible_with(t)) if expected_shape is None: self.assertIsNone(result.shape.rank) else: self.assertEqual(result.shape.as_list(), expected_shape) def test_with_keras_tensor_with_ragged_spec(self): t = keras.engine.keras_tensor.KerasTensor( tf.RaggedTensorSpec(shape=(None, None, 1))) self.assertIsInstance(tf_utils.get_tensor_spec(t), tf.RaggedTensorSpec)
def setUp(self): super().setUp() self._data_context = data_converter.DataContext( time_step_spec=ts.TimeStep(step_type=(), reward=tf.TensorSpec((), tf.float32), discount=tf.TensorSpec((), tf.float32), observation=()), action_spec={'action1': tf.TensorSpec((), tf.float32)}, policy_state_spec=(), info_spec=()) self._data_context_with_state = data_converter.DataContext( time_step_spec=ts.TimeStep(step_type=(), reward=tf.TensorSpec((), tf.float32), discount=tf.TensorSpec((), tf.float32), observation=tf.TensorSpec((2, ), tf.float32)), action_spec={'action1': tf.TensorSpec((), tf.float32)}, policy_state_spec=[ tf.TensorSpec((2, ), tf.float32), tf.TensorSpec((2, ), tf.float32) ], info_spec=())
class TensorListModule(tf.Module): def __init__(self): pass @tf.function(input_signature=[tf.TensorSpec([], tf.float32)]) def identity_through_tensorlist(self, x): ta = tf.TensorArray(dtype=tf.float32, size=1, element_shape=[]) ta = ta.write(0, x) return ta.read(0) @tf.function(input_signature=[ tf.TensorSpec([], tf.float32), tf.TensorSpec([], tf.float32) ]) def add_through_tensorlist(self, a, b): ta = tf.TensorArray(dtype=tf.float32, size=2, element_shape=[]) ta = ta.write(0, a) ta = ta.write(1, b) return ta.read(0) + ta.read(1) @tf.function(input_signature=[tf.TensorSpec([STATIC_SIZE], tf.float32)]) def slice_first_element_with_from_tensor(self, t): ta = tf.TensorArray(dtype=tf.float32, size=STATIC_SIZE, element_shape=[]) ta = ta.unstack(t) return ta.read(0) @tf.function(input_signature=[ tf.TensorSpec([STATIC_SIZE, STATIC_SIZE], tf.float32) ]) def slice_first_element_with_from_tensor_high_rank(self, t): ta = tf.TensorArray(dtype=tf.float32, size=STATIC_SIZE, element_shape=[STATIC_SIZE]) ta = ta.unstack(t) return ta.read(0) @tf.function(input_signature=[ tf.TensorSpec([], tf.float32), tf.TensorSpec([], tf.float32) ]) def concat_with_tensorlist_stack(self, a, b): ta = tf.TensorArray(dtype=tf.float32, size=2, element_shape=[]) ta = ta.write(0, a) ta = ta.write(1, b) return ta.stack()
def create_serving_signature(self, model, feature_mapper, label_inverse_lookup_layer): """Create serving signature for the given model.""" @tf.function def serve_fn(raw_features): raw_features = tf.compat.v1.expand_dims(raw_features, axis=0) transformed_features = model.feature_mapper(raw_features) outputs = model(transformed_features) outputs = tf.compat.v1.squeeze(outputs, axis=0) outputs = tf.cast(tf.greater(outputs, 0.5), tf.int64) decoded_outputs = model.label_inverse_lookup_layer(outputs) return tf.compat.v1.squeeze(decoded_outputs, axis=0) model.feature_mapper = feature_mapper model.label_inverse_lookup_layer = label_inverse_lookup_layer # serving does NOT have batch dimension return serve_fn.get_concrete_function( tf.TensorSpec(shape=(3), dtype=tf.string, name="example"))
def test_conv1d_recreate_conv_unknown_dims(self): with self.cached_session(): layer = keras.layers.Conv1D(filters=1, kernel_size=3, strides=1, dilation_rate=2, padding='causal') inpt1 = np.random.normal(size=[1, 9, 1]).astype(np.float32) inpt2 = np.random.normal(size=[1, 2, 1]).astype(np.float32) outp1_shape = layer(inpt1).shape @tf.function(input_signature=[tf.TensorSpec([1, None, 1])]) def fn(inpt): return layer(inpt) fn(inpt2) self.assertEqual(outp1_shape, layer(inpt1).shape)
def create_reverb_server_for_replay_buffer_and_variable_container( collect_policy, train_step, replay_buffer_capacity, port): """Sets up one reverb server for replay buffer and variable container.""" # Create the signature for the variable container holding the policy weights. variables = { reverb_variable_container.POLICY_KEY: collect_policy.variables(), reverb_variable_container.TRAIN_STEP_KEY: train_step } variable_container_signature = tf.nest.map_structure( lambda variable: tf.TensorSpec(variable.shape, dtype=variable.dtype), variables) # Create the signature for the replay buffer holding observed experience. replay_buffer_signature = tensor_spec.from_spec( collect_policy.collect_data_spec) # TODO(b/188427258) Add time dimension when using Reverb.TrajectoryWriters. # replay_buffer_signature = tensor_spec.add_outer_dim(replay_buffer_signature) # Crete and start the replay buffer and variable container server. server = reverb.Server( tables=[ reverb.Table( # Replay buffer storing experience. name=reverb_replay_buffer.DEFAULT_TABLE, sampler=reverb.selectors.Uniform(), remover=reverb.selectors.Fifo(), # TODO(b/159073060): Set rate limiter for SAC properly. rate_limiter=reverb.rate_limiters.MinSize(1), max_size=replay_buffer_capacity, max_times_sampled=0, signature=replay_buffer_signature, ), reverb.Table( # Variable container storing policy parameters. name=reverb_variable_container.DEFAULT_TABLE, sampler=reverb.selectors.Uniform(), remover=reverb.selectors.Fifo(), rate_limiter=reverb.rate_limiters.MinSize(1), max_size=1, max_times_sampled=0, signature=variable_container_signature, ), ], port=port) return server
def __init__(self): super().__init__() self.m = initialize_model() input_shape = list([BATCH_SIZE] + self.m.inputs[0].shape[1:]) # Some models accept dynamic image dimensions by default, so we use # IMAGE_DIM as a stand-in. for i, dim in enumerate(input_shape): if dim is None: input_shape[i] = IMAGE_DIM # Specify input shape with a static batch size. # TODO(b/142948097): Add support for dynamic shapes in SPIR-V lowering. self.call = tf_test_utils.tf_function_unit_test( input_signature=[tf.TensorSpec(input_shape)], name="call", rtol=1e-5, atol=1e-5)(lambda x: self.m(x, training=False))
def get_model(): relay_file = "relay.json" relay_params = "relay.params" from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 model = tf.keras.Sequential([ hub.KerasLayer(tf_hub_links[model_name], output_shape=[1001]) ]) img_size = 299 if model_name == 'inceptionv3' else 224 np_image = np.random.rand(1, img_size, img_size, 3).astype('float32') model._set_inputs(np_image) # Convert Keras model to ConcreteFunction full_model = tf.function(lambda x: model(x)) full_model = full_model.get_concrete_function( tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype, name="data")) frozen_func = convert_variables_to_constants_v2(full_model) frozen_func.graph.as_graph_def() tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir="./.tf_saved_model/" + model_name, name="frozen_graph.pb", as_text=False) parser = tvm.relay.frontend.TFParser("./.tf_saved_model/" + model_name + "/frozen_graph.pb") graph_def = parser.parse() mod, params = relay.frontend.from_tensorflow(graph_def, shape={"data": (1, img_size, img_size, 3)}) # with open(relay_file, "w") as fo: # fo.write(tvm.ir.save_json(mod)) # with open(relay_params, "wb") as fo: # fo.write(relay.save_param_dict(params)) # with open(relay_file, "r") as fi: # mod = tvm.ir.load_json(fi.read()) # with open(relay_params, "rb") as fi: # params = relay.load_param_dict(fi.read()) return mod, params
def testKerasModel(self): """Test a basic Keras model with Variables.""" input_data = {"x": tf.constant(1., shape=[1, 1])} # Create a simple Keras model. x = [-1, 0, 1, 2, 3, 4] y = [-3, -1, 1, 3, 5, 7] model = keras.models.Sequential( [keras.layers.Dense(units=1, input_shape=[1])]) model.compile(optimizer="sgd", loss="mean_squared_error") model.fit(x, y, epochs=1) @tf.function( input_signature=[tf.TensorSpec(shape=[1, 1], dtype=tf.float32)]) def to_save(x): return model(x) root, output_func = self._freezeModel(to_save) self._testConvertedFunction(root, root.f, output_func, input_data)
def testKerasLSTM(self): """Test a Keras LSTM containing dynamic_rnn ops.""" input_data = { "x": tf.constant( np.array(np.random.random_sample((10, 10, 10)), dtype=np.float32)) } model = keras.models.Sequential( [keras.layers.LSTM(units=10, input_shape=(10, 10))]) @tf.function(input_signature=[ tf.TensorSpec(shape=[10, 10, 10], dtype=tf.float32) ]) def to_save(x): return model(x) root, output_func = self._freezeModel(to_save) self._testConvertedFunction(root, root.f, output_func, input_data)
def build_saved_model(model): """Returns a tf.Module for saving to SavedModel.""" class SimCLRModel(tf.Module): """Saved model for exporting to hub.""" def __init__(self, model): self.model = model # This can't be called `trainable_variables` because `tf.Module` has # a getter with the same name. self.trainable_variables_list = model.trainable_variables @tf.function def __call__(self, inputs, trainable): self.model(inputs, training=trainable) return get_salient_tensors_dict() module = SimCLRModel(model) input_spec = tf.TensorSpec(shape=[None, None, None, 3], dtype=tf.float32) module.__call__.get_concrete_function(input_spec, trainable=True) module.__call__.get_concrete_function(input_spec, trainable=False) return module
def models(): tf.keras.backend.set_learning_phase(False) # 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
def __init__(self, unroll_length=1): self._env = MockEnv(state_space_size=4, unroll_length=unroll_length) self._agent = MockAgent(unroll_length=unroll_length) self._actor_output_spec = common.ActorOutput( initial_agent_state=tf.TensorSpec(shape=[5], dtype=tf.float32), env_output=self._env.env_spec, agent_output=self._agent.agent_spec, actor_action=common.ActorAction( chosen_action_idx=tf.TensorSpec(shape=[unroll_length + 1], dtype=tf.int32), oracle_next_action_idx=tf.TensorSpec(shape=[unroll_length + 1], dtype=tf.int32), action_val=tf.TensorSpec(shape=[unroll_length + 1], dtype=tf.int32), log_prob=tf.TensorSpec(shape=[unroll_length + 1], dtype=tf.float32)), loss_type=tf.TensorSpec(shape=[], dtype=tf.int32), info=tf.TensorSpec(shape=[], dtype=tf.string), )
def test_run_eval_actor_once(self): hparams = {} hparams['max_iter'] = 1 hparams['num_episodes_per_iter'] = 5 hparams['logdir'] = os.path.join(FLAGS.test_tmpdir, 'model') mock_problem = testing_utils.MockProblem( unroll_length=FLAGS.unroll_length) agent = mock_problem.get_agent() ckpt_manager = _get_ckpt_manager(hparams['logdir'], agent=agent) checkpoint_path = ckpt_manager.save(checkpoint_number=0) # Create a no-op gRPC server that responds to Aggregator RPCs. server_address = 'unix:/tmp/eval_actor_test_grpc' server = grpc.Server([server_address]) @tf.function( input_signature=[tf.TensorSpec(shape=(), dtype=tf.string)]) def eval_enqueue(_): return [] # Test 01. Eval with aggregator. server.bind(eval_enqueue) server.start() eval_actor.run_evaluation(mock_problem, server_address, hparams) # Test 02. Eval without aggregator. hparams['task_id'] = 0000 hparams['max_iter'] = 1 eval_actor.run_evaluation(mock_problem, None, hparams, checkpoint_path, FLAGS.test_tmpdir, file_prefix='mock_test', test_mode=True) with gfile.GFile(os.path.join(FLAGS.test_tmpdir, 'test_data_0.p'), 'rb') as fp: eval_result = pickle.load(fp) self.assertEqual(eval_result['mock_test_0_0_0000']['result'], 1000.0)
class MNIST(tf.keras.models.Model): """Model representing a MNIST classifier.""" def __init__(self, output_activation="softmax"): super(MNIST, self).__init__() self.layer_1 = tf.keras.layers.Dense(64) self.layer_2 = tf.keras.layers.Dense(10, activation=output_activation) @tf.function(input_signature=[ tf.TensorSpec(shape=[None, 28, 28, 1], dtype=tf.uint8) ]) def call(self, inputs): casted = tf.keras.layers.Lambda(lambda x: tf.cast(x, tf.float32))( inputs) flatten = tf.keras.layers.Flatten()(casted) def normalize_fn(x): return x / tf.reduce_max(tf.gather(x, 0)) normalize = tf.keras.layers.Lambda(normalize_fn)(flatten) x = self.layer_1(normalize) output = self.layer_2(x) return output
def testReproduceBug159550941(self): # Reproduction for b/159550941. input_signature = [tf.TensorSpec([], tf.int32)] @tf.function(input_signature=input_signature) def sample(chains): initial_state = tf.zeros([chains, 1]) def log_prob(x): return tf.reduce_sum(tfp.distributions.Normal(0, 1).log_prob(x), -1) kernel = tfp.mcmc.HamiltonianMonteCarlo( target_log_prob_fn=log_prob, num_leapfrog_steps=3, step_size=1e-3) return tfp.mcmc.sample_chain( num_results=5, num_burnin_steps=4, current_state=initial_state, kernel=kernel, trace_fn=None) # Checking that shape inference doesn't fail. sample(2)
def models(): tf.keras.backend.set_learning_phase(False) tf_test_utils.set_random_seed() input_shape = get_input_shape(FLAGS.data, FLAGS.model) # 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 # if weights == 'imagenet' it will load weights from external tf.keras URL model = APP_MODELS[FLAGS.model]( weights=weights, include_top=FLAGS.include_top, 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, 'cifar10_include_top_{}_{}'.format(FLAGS.include_top, FLAGS.model + '.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