def test_tf_saved_model_save_multiple_signatures(self):
base_path = os.path.join(self.get_temp_dir(), 'tf_saved_model_save')
export_path = os.path.join(base_path, '00000123')
root = tf.train.Checkpoint()
root.f = tf.function(lambda x: {'y': 1.},
input_signature=[tf.TensorSpec(None, tf.float32)])
root.g = tf.function(lambda x: {'y': 2.},
input_signature=[tf.TensorSpec(None, tf.float32)])
tf.saved_model.experimental.save(
root,
export_path,
signatures={
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: root.f,
'custom_signature_key': root.g
})
_, model_server_address, _ = TensorflowModelServerTest.RunServer(
'default', base_path)
expected_version = self._GetModelVersion(base_path)
self.VerifyPredictRequest(model_server_address,
expected_output=2.0,
expected_version=expected_version,
signature_name='custom_signature_key')
self.VerifyPredictRequest(model_server_address,
expected_output=1.0,
expected_version=expected_version)
def training_loop(train_dataset,
valid_dataset,
model,
hparams):
"""Trains a GNP for a fixed number of iterations."""
optimizer_config = {'optimizer': hparams.optimizer(hparams.lr),
'max_grad_norm': hparams.max_grad_norm}
num_context = hparams.num_context
best_recon_loss = np.inf
if hparams.is_nll:
step = tf.function(utils.nll_gnp_step_bandits.python_function)
valid_metric = utils.nll
else:
step = tf.function(utils.mse_gnp_step_bandits.python_function)
valid_metric = utils.mse
for it in range(hparams.num_iterations):
batch_train_data = get_splits(
train_dataset,
num_context,
hparams.batch_size,
points_perm=True)
recon_loss, local_z_kl, global_z_kl = step(
model,
batch_train_data,
optimizer_config)
if it % hparams.print_every == 0:
batch_valid_data = get_splits(
valid_dataset,
num_context,
hparams.batch_size,
points_perm=False)
(batch_context_x,
batch_context_y,
batch_target_x,
batch_target_y,
batch_unseen_target_y,
batch_unseen_target_a) = batch_valid_data
prediction = model(batch_context_x,
batch_context_y,
batch_target_x,
batch_target_y)
batch_unseen_predictions = prediction[:, num_context:]
valid_recon_loss = valid_metric(batch_unseen_target_y,
batch_unseen_predictions,
batch_unseen_target_a)
print('it: {}, train recon loss: {}, local kl: {} global kl: {} '
'valid reconstr loss: {}'.format(it, recon_loss, local_z_kl,
global_z_kl, valid_recon_loss))
if valid_recon_loss.numpy() < best_recon_loss:
best_recon_loss = valid_recon_loss.numpy()
print('Saving best model with reconstruction loss',
best_recon_loss, flush=True)
model.save_weights(hparams.save_path)
def testNormalNormalSample(self):
# Standard normal prior.
# Samples are shape [2].
normal_prior = tfd.Normal(self.dtype([0., 0.]), self.dtype(1.))
def normal_sampler(seed):
return normal_prior.sample(seed=seed)
# A single data point at the mode.
# The state is expected to be 2 dimensional, so
# we reduce sum on the last axis.
def normal_log_likelihood(state):
return tf.reduce_sum(tfd.Normal(state, self.dtype(2.)).log_prob(
self.dtype(0.)),
axis=-1)
kernel = elliptical.EllipticalSliceSampler(
normal_sampler_fn=normal_sampler,
log_likelihood_fn=normal_log_likelihood,
seed=tfp_test_util.test_seed(),
)
samples = tf.function(lambda: tfp.mcmc.sample_chain( # pylint: disable=g-long-lambda
num_results=int(3e5),
current_state=self.dtype(np.random.randn(2)),
kernel=kernel,
num_burnin_steps=int(1e4),
parallel_iterations=1,
trace_fn=None))()
mean, variance = self.evaluate(tf.nn.moments(samples, axes=[0]))
# Computed exactly from the formula in normal-normal posterior.
self.assertAllClose([0., 0.], mean, rtol=5e-2, atol=6e-3)
self.assertAllClose([4. / 5, 4. / 5], variance, rtol=5e-2)
def graph_fn():
boxes = tf.zeros([0, 4], dtype=tf.float32)
blackout = tf.zeros([0], dtype=tf.bool)
blackout_pixel_weights_by_box_regions = tf.function(
ta_utils.blackout_pixel_weights_by_box_regions)
output = blackout_pixel_weights_by_box_regions(
10, 20, boxes, blackout)
return output
def graph_fn():
boxes = tf.constant(
[[0.0, 0.0, 5, 5], [0.0, 0.0, 10.0, 20.0], [6.0, 12.0, 8.0, 18.0]],
dtype=tf.float32)
blackout = tf.constant([True, False, True], dtype=tf.bool)
blackout_pixel_weights_by_box_regions = tf.function(
ta_utils.blackout_pixel_weights_by_box_regions)
output = blackout_pixel_weights_by_box_regions(10, 20, boxes, blackout)
return output
def test_linear(self, solver):
jacobian_diag_part = np.float32([-0.5, -1.])
initial_state = np.float32([1., 2.])
fn = lambda: linear(solver, jacobian_diag_part, initial_state)
fn = tf.function(fn, autograph=False, jit_compile=True)
with tf.device(FLAGS.test_device):
times, states = self.evaluate(fn())
states_exact = np.exp(jacobian_diag_part[np.newaxis, :] *
times[:, np.newaxis]) * initial_state
self.assertAllClose(states, states_exact, rtol=1e-4)
def test_hub_model(self):
image = tf.random.uniform((1, 320, 320, 3))
keras_model = efficientdet_keras.EfficientDetNet('efficientdet-lite0')
tmp_ckpt = os.path.join(tempfile.mkdtemp(), 'ckpt')
keras_model.config.model_dir = tmp_ckpt
base_model = train_lib.EfficientDetNetTrainHub(
keras_model.config,
"https://tfhub.dev/tensorflow/efficientdet/lite0/feature-vector/1")
cls_outputs, box_outputs = tf.function(base_model)(image,
training=False)
keras_model.build(image.shape)
d1 = {var.name: var for var in base_model.variables}
for var in keras_model.variables:
var.assign(d1[var.name].numpy())
cls_outputs2, box_outputs2 = tf.function(keras_model)(image, False)
for c1, b1, c2, b2 in zip(cls_outputs, box_outputs, cls_outputs2,
box_outputs2):
self.assertAllEqual(c1, c2)
self.assertAllEqual(b1, b2)
def graph_fn():
boxes = tf.constant([[0.0, 0.0, 2.0, 2.0], [0.0, 0.0, 4.0, 2.0],
[3.0, 0.0, 4.0, 4.0]],
dtype=tf.float32)
blackout = tf.constant([False, False, True], dtype=tf.bool)
weights = tf.constant([0.4, 0.3, 0.2], tf.float32)
blackout_pixel_weights_by_box_regions = tf.function(
ta_utils.blackout_pixel_weights_by_box_regions)
output = blackout_pixel_weights_by_box_regions(
4, 4, boxes, blackout, weights)
return output
def __init__(self, name, hparams, optimizer='RMS'):
self.name = name
self.hparams = hparams
self.verbose = getattr(hparams, 'verbose', True)
self.update_freq_lr = hparams.training_freq
self.update_freq_nn = hparams.training_freq_network
self.t = 0
self.num_epochs = hparams.training_epochs
self.data_h = contextual_dataset.ContextualDataset(hparams.context_dim,
hparams.num_actions,
intercept=False)
self.gradient_updates = tf.Variable(0, trainable=False)
if self.hparams.activate_decay:
self.lr = tf.train.inverse_time_decay(self.hparams.initial_lr,
self.gradient_updates, 1,
self.hparams.lr_decay_rate)
else:
self.lr = tf.Variable(self.hparams.initial_lr, trainable=False)
optimizer = tf.train.RMSPropOptimizer(self.lr)
self._optimizer_config = {
'optimizer': optimizer,
'max_grad_norm': hparams.max_grad_norm
}
if self.verbose:
print('Initializing model {}.'.format(self.name))
self.snp = regressor.Regressor(
input_dim=hparams.context_dim + hparams.num_actions,
output_dim=1,
x_encoder_sizes=hparams.x_encoder_sizes,
x_y_encoder_sizes=hparams.x_y_encoder_sizes,
global_latent_net_sizes=hparams.global_latent_net_sizes,
local_latent_net_sizes=hparams.local_latent_net_sizes,
heteroskedastic_net_sizes=hparams.heteroskedastic_net_sizes,
att_type=hparams.att_type,
att_heads=hparams.att_heads,
uncertainty_type=hparams.uncertainty_type,
mean_att_type=hparams.mean_att_type,
scale_att_type_1=hparams.scale_att_type_1,
scale_att_type_2=hparams.scale_att_type_2,
activation=hparams.activation,
output_activation=hparams.output_activation,
data_uncertainty=hparams.data_uncertainty,
local_variational=hparams.local_variational,
model_path=hparams.model_path)
self._step = tf.function(utils.mse_step.python_function) # pytype: disable=module-attr
self._one_hot_vectors = tf.one_hot(indices=np.arange(
hparams.num_actions),
depth=hparams.num_actions)
def training_loop(train_dataset,
valid_dataset,
model,
hparams):
"""Trains an SNP for a fixed number of iterations."""
optimizer_config = {'optimizer': hparams.optimizer(hparams.lr),
'max_grad_norm': hparams.max_grad_norm}
num_context = hparams.num_context
best_mse = np.inf
step = tf.function(utils.mse_step.python_function) # pytype: disable=module-attr
for it in range(hparams.num_iterations):
batch_train_data = get_splits(
train_dataset,
num_context,
hparams.batch_size,
points_perm=True)
nll, mse, local_z_kl, global_z_kl = step(
model,
batch_train_data,
optimizer_config)
if it % hparams.print_every == 0:
(batch_context_x,
batch_context_y,
batch_target_x,
batch_target_y,
batch_unseen_targets) = get_splits(valid_dataset,
num_context,
hparams.batch_size,
points_perm=False)
prediction = model(batch_context_x,
batch_context_y,
batch_target_x,
batch_target_y)
batch_unseen_predictions = prediction[:, num_context:]
valid_nll = utils.nll(batch_unseen_targets, batch_unseen_predictions)
valid_mse = utils.mse(batch_unseen_targets, batch_unseen_predictions)
if model.local_variational:
valid_local_kl = tf.reduce_mean(
tf.reduce_sum(model.losses[-1][:, num_context:], axis=[1, 2]))
else:
valid_local_kl = 0.
valid_global_kl = tf.reduce_mean(tf.reduce_sum(model.losses[-2], axis=-1))
print('it: {}, train nll: {}, mse: {}, local kl: {} global kl: {} '
'valid nll: {}, mse: {}, local kl: {} global kl: {}'
.format(it, nll, mse, local_z_kl, global_z_kl,
valid_nll, valid_mse, valid_local_kl, valid_global_kl))
if valid_mse.numpy() < best_mse:
best_mse = valid_mse.numpy()
print('Saving best model with MSE', best_mse)
model.save_weights(hparams.save_path)
def testSampleChainSeedReproducible(self):
normal_prior = tfd.Normal(5 * [[0., 0.]], 1.)
def normal_sampler(seed):
return normal_prior.sample(seed=seed)
def normal_log_likelihood(state):
return tf.reduce_sum(tfd.Normal(state, 2.).log_prob(0.), axis=-1)
num_results = 10
seed = tfp_test_util.test_seed()
current_state = np.float32(np.random.rand(5, 2))
samples0 = tf.function(lambda: tfp.mcmc.sample_chain( # pylint: disable=g-long-lambda
num_results=2 * num_results,
num_steps_between_results=0,
current_state=current_state,
kernel=elliptical.EllipticalSliceSampler(
normal_sampler_fn=normal_sampler,
log_likelihood_fn=normal_log_likelihood,
seed=seed),
num_burnin_steps=150,
trace_fn=None,
parallel_iterations=1))()
samples1 = tf.function(lambda: tfp.mcmc.sample_chain( # pylint: disable=g-long-lambda
num_results=num_results,
num_steps_between_results=1,
current_state=current_state,
kernel=elliptical.EllipticalSliceSampler(
normal_sampler_fn=normal_sampler,
log_likelihood_fn=normal_log_likelihood,
seed=seed),
trace_fn=None,
num_burnin_steps=150,
parallel_iterations=1))()
samples0_, samples1_ = self.evaluate([samples0, samples1])
self.assertAllClose(samples0_[::2], samples1_, atol=1e-5, rtol=1e-5)
def testMultivariateNormalNd(self, event_size=32, batch_size=8, num_steps=2):
tf.set_random_seed(3)
with tf.device(FLAGS.test_device):
f = run_nuts_chain(event_size, batch_size, num_steps)
f = tf.function(f, autograph=False, jit_compile=True)
samples, leapfrogs = self.evaluate(f())
# TODO(axch) Figure out what the right thing to test about the leapfrog
# count really is and test it, instead of just flailing around like this
# does.
print(type(samples), type(leapfrogs))
print(samples, leapfrogs)
ev_leapfrogs = leapfrogs[0]
self.assertGreater(len(set(ev_leapfrogs.tolist())), 1)
self.assertTrue(all(ev_leapfrogs > 1))
def test_tf_saved_model_save(self):
base_path = os.path.join(self.get_temp_dir(), 'tf_saved_model_save')
export_path = os.path.join(base_path, '00000123')
root = tf.train.Checkpoint()
root.v1 = tf.Variable(3.)
root.v2 = tf.Variable(2.)
root.f = tf.function(lambda x: {'y': root.v1 * root.v2 * x})
to_save = root.f.get_concrete_function(tf.TensorSpec(None, tf.float32))
tf.saved_model.experimental.save(root, export_path, to_save)
_, model_server_address, _ = TensorflowModelServerTest.RunServer(
'default', base_path)
expected_version = self._GetModelVersion(base_path)
self.VerifyPredictRequest(model_server_address,
expected_output=12.0,
specify_output=False,
expected_version=expected_version)
def _tfTestHelper(self, run_asserts_fn, execute_program_fn):
# Note: test_device is 'cpu', 'gpu', etc.
# Various int32 and int64 kernels are missing for GPU, so we skip direct
# tests on the GPU device, but test XLA on GPU below.
if 'cpu' in FLAGS.test_device.lower():
# Make sure everything works with no XLA compilation.
with tf.device('CPU:0'):
run_asserts_fn(
functools.partial(execute_program_fn, backend=TF_BACKEND))
# Force XLA compilation using tf.function.
backend = TF_BACKEND_NO_ASSERTS
f = functools.partial(execute_program_fn, backend=backend)
f = tf.function(f, autograph=False, experimental_compile=True)
with tf.device(FLAGS.test_device):
run_asserts_fn(f)
def testNormalNormalSampleMultipleDatapoints(self):
# Two independent chains, of states of shape [3].
prior_stddev = self.dtype(np.exp(np.random.rand(2, 3)))
likelihood_stddev = self.dtype(np.exp(np.random.rand(2, 3)))
# 10 data points.
data = self.dtype(np.random.randn(10, 2, 3))
# Standard normal prior.
normal_prior = tfd.Normal(self.dtype(0.), prior_stddev)
def normal_sampler(seed):
return normal_prior.sample(seed=seed)
# 10 samples at 2 chains.
def normal_log_likelihood(state):
return tf.reduce_sum(
tfd.Normal(state, likelihood_stddev).log_prob(data),
axis=[0, -1],
)
kernel = elliptical.EllipticalSliceSampler(
normal_sampler_fn=normal_sampler,
log_likelihood_fn=normal_log_likelihood,
seed=tfp_test_util.test_seed(),
)
samples = tf.function(lambda: tfp.mcmc.sample_chain( # pylint: disable=g-long-lambda
num_results=int(3e5),
current_state=self.dtype(np.random.randn(2, 3)),
kernel=kernel,
num_burnin_steps=int(1e4),
parallel_iterations=1,
trace_fn=None))()
mean, variance = self.evaluate(tf.nn.moments(samples, axes=[0]))
posterior_mean, posterior_variance = normal_normal_posterior(
prior_mean=0.,
prior_stddev=prior_stddev,
likelihood_stddev=likelihood_stddev,
data=data)
# Computed exactly from the formula in normal-normal posterior.
self.assertAllClose(posterior_mean, mean, rtol=2e-2, atol=6e-3)
self.assertAllClose(posterior_variance, variance, rtol=5e-2)
def __init__(self, weight_path):
helpers.ensure_lpips_weights_exist(weight_path)
def wrap_frozen_graph(graph_def, inputs, outputs):
def _imports_graph_def():
tf.graph_util.import_graph_def(graph_def, name="")
wrapped_import = tf.wrap_function(_imports_graph_def, [])
import_graph = wrapped_import.graph
return wrapped_import.prune(
tf.nest.map_structure(import_graph.as_graph_element, inputs),
tf.nest.map_structure(import_graph.as_graph_element, outputs))
# Pack LPIPS network into a tf function
graph_def = tf.GraphDef()
with open(weight_path, "rb") as f:
graph_def.ParseFromString(f.read())
self._lpips_func = tf.function(
wrap_frozen_graph(
graph_def, inputs=("0:0", "1:0"), outputs="Reshape_10:0"))
def reset_tf_function(tf_function):
return tf.function(tf_function.python_function)
def reset(self):
self.tf_function = tf.function(self.python_function)
def __init__(self, python_function):
self.tf_function = tf.function(python_function)
self.python_function = python_function
def multiply2n_ragged(tensor1, tensor2):
#this function multiplies two ragged tesnsors of rank 2 . the most outer ranks of the two tensros must be equal .
#setting variables and constats
outerloop_counter = tf.constant(0, dtype=tf.int32)
carry_on = tf.constant(0, dtype=tf.int32)
taValues = tf.TensorArray(tf.float32,
size=0,
dynamic_size=True,
clear_after_read=False,
infer_shape=False)
taL2Splits = tf.TensorArray(tf.int32,
size=0,
dynamic_size=True,
clear_after_read=False,
infer_shape=False)
taL1Splits = tf.TensorArray(tf.int32,
size=0,
dynamic_size=True,
clear_after_read=False,
infer_shape=False)
taL1Splits = taL1Splits.write(
0, [0]) ## required intialization for L1 split only
innerloop_processing_graphed = tf.function(innerloop_processing)
generateL1Tensor_writeback_graphed = tf.function(
generateL1Tensor_writeback)
def outerloop_cond(counter, input1, input2, taValues, taL2Splits,
taL1Splits, carry_on):
value = tf.shape(input1[2])[0] - 1
return counter < value ## this is the length of the outermost dimision , stop of this
def outloop_body(counter, input1, input2, taValues, taL2Splits, taL1Splits,
carry_on):
l1_comp_begin = input1[2][
counter] ## this is begin position of the current row in the outer split ( ie. the ith value in the outer row split tensor )
l1_comp_end = input1[2][
counter +
1] ## this is end position of the current row in the outer split (ie. the ith + 1 value in the outer row split tensor)
l1_comp2_begin = input2[2][
counter] ## we do the same for the second components
l1_comp2_end = input2[2][
counter + 1] ## we do the same for the second components
comp = innerloop_processing_graphed(
l1_comp_begin, l1_comp_end, input1
) ## now retrive the data to be procesed for the selected rows from vector1
comp2 = innerloop_processing_graphed(
l1_comp2_begin, l1_comp2_end, input2) ## do the same for vector 2
#comp2 = tf.transpose(comp2) ### desired operation
multiply = tf.matmul(comp, comp2) #### This is the desired operation
myshape = tf.shape(
multiply
) ## calculate the shape of the result in order to prepare to write the result in a ragged tensor format.
offset = tf.cond(
taValues.size() > 0, lambda: tf.shape(taValues.concat())[0],
lambda: 0
) ### this is a hack, TensorArray.concat returns an error if the array is empty. Thus we check before calling this.
#print11=tf.print("=================Final Shape is : " ,myshape[0] , " X " ,myshape[1] )
l2v = generateL1Tensor_writeback_graphed(
offset, myshape[1], myshape[0]
) # generate the inner row split of the result for the current element
taL2Splits = taL2Splits.write(
counter, l2v) # write back the inner rowlplit to a TensorArray
taValues = taValues.write(
counter, tf.reshape(multiply, [-1])
) # wirte back the actual ragged tensor elemnts in a another TensorArray
carry_on = carry_on + myshape[
0] ## required to calculate the outer row splite
taL1Splits = taL1Splits.write(
counter + 1, [carry_on]) ## This is the outmost row split.
with tf.control_dependencies(
[comp, comp2, myshape, l2v, carry_on, multiply]):
counter = counter + 1
return counter, input1, input2, taValues, taL2Splits, taL1Splits, carry_on
with tf.name_scope("RaggedMultiply"):
outerloop_finalcounter, _, _, ta1, ta2, ta3, _ = tf.while_loop(
outerloop_cond,
outloop_body, [
outerloop_counter, tensor1, tensor2, taValues, taL2Splits,
taL1Splits, carry_on
],
back_prop=True)
uinquie_ta2, _ = tf.unique(
ta2.concat()
) # this is required since some values might be duplicate in the row split itself
t1 = ta1.concat()
t3 = ta3.concat()
#with tf.control_dependencies([t1 , uinquie_ta2 ,t3 ]):
final_values = t1, uinquie_ta2, t3
return final_values
