def convolute_and_save(module_path,
signature,
export_path,
transform_fn,
transform_checkpoint_path,
new_signature=None):
"""Loads TFHub module, convolutes it with transform_fn and saves it again.
Args:
module_path: String with path from which the module is constructed.
signature: String with name of signature to use for loaded module.
export_path: String with path where to save the final TFHub module.
transform_fn: Function that creates the graph to be appended to the loaded
TFHub module. The function should take as keyword arguments the tensors
returned by the loaded TFHub module. The function should return a
dictionary of tensor that will be the output of the new TFHub module.
transform_checkpoint_path: Path to checkpoint from which the transformer_fn
variables will be read.
new_signature: String with new name of signature to use for saved module. If
None, `signature` is used instead.
"""
if new_signature is None:
new_signature = signature
# We create a module_fn that creates the new TFHub module.
def module_fn():
module = hub.Module(module_path)
inputs = _placeholders_from_module(module, signature=signature)
intermediate_tensor = module(inputs, signature=signature, as_dict=True)
# We need to scope the variables that are created when the transform_fn is
# applied.
with tf.variable_scope("transform"):
outputs = transform_fn(**intermediate_tensor)
hub.add_signature(name=new_signature, inputs=inputs, outputs=outputs)
# We create a new graph where we will build the module for export.
with tf.Graph().as_default():
# Create the module_spec for the export.
spec = hub.create_module_spec(module_fn)
m = hub.Module(spec, trainable=True)
# We need to recover the scoped variables and remove the scope when loading
# from the checkpoint.
prefix = "transform/"
transform_variables = {
k[len(prefix):]: v
for k, v in m.variable_map.items() if k.startswith(prefix)
}
if transform_variables:
init_fn = tf.contrib.framework.assign_from_checkpoint_fn(
transform_checkpoint_path, transform_variables)
with tf.Session() as sess:
# Initialize all variables, this also loads the TFHub parameters.
sess.run(tf.global_variables_initializer())
# Load the transformer variables from the checkpoint.
if transform_variables:
init_fn(sess)
# Export the new TFHub module.
m.export(export_path, sess)
def testClearControlDependenciesForModuleStateButNotApplyGraphs(self):
module_spec = hub.create_module_spec(stateless_module_fn)
with tf.Graph().as_default() as g1:
v = tf.placeholder(dtype=tf.int64, name="v")
m = hub.Module(module_spec)
m(v)
with tf.Graph().as_default() as g2:
v = tf.placeholder(dtype=tf.int64, name="v")
with tf.control_dependencies([v]):
m = hub.Module(module_spec)
m(v)
self.assertEqual(g1.as_graph_def(), g2.as_graph_def())
with tf.Graph().as_default() as g3:
v = tf.placeholder(dtype=tf.int64, name="v")
m = hub.Module(module_spec)
m(v)
with tf.Graph().as_default() as g4:
v = tf.placeholder(dtype=tf.int64, name="v")
m = hub.Module(module_spec)
with tf.control_dependencies([v]):
m(v)
self.assertNotEqual(g3.as_graph_def(), g4.as_graph_def())
def test_module_export_vocab_on_custom_fs(self):
root_dir = "file://%s" % self.get_temp_dir()
export_dir = "%s_%s" % (root_dir, "export")
tf.gfile.MakeDirs(export_dir)
# Create a module with a vocab file located on a custom filesystem.
vocab_dir = os.path.join(root_dir, "vocab_location")
tf.gfile.MakeDirs(vocab_dir)
vocab_filename = os.path.join(vocab_dir, "tokens.txt")
tf_utils.atomic_write_string_to_file(vocab_filename, "one", False)
def create_assets_module_fn():
def assets_module_fn():
indices = tf.placeholder(dtype=tf.int64, name="indices")
table = tf.contrib.lookup.index_to_string_table_from_file(
vocabulary_file=vocab_filename, default_value="UNKNOWN")
outputs = table.lookup(indices)
hub.add_signature(inputs=indices, outputs=outputs)
return assets_module_fn
with tf.Graph().as_default():
assets_module_fn = create_assets_module_fn()
spec = hub.create_module_spec(assets_module_fn)
embedding_module = hub.Module(spec)
with tf.Session() as sess:
sess.run(tf.tables_initializer())
embedding_module.export(export_dir, sess)
module_files = tf.gfile.ListDirectory(export_dir)
self.assertListEqual(
["assets", "saved_model.pb", "tfhub_module.pb", "variables"],
sorted(module_files))
module_files = tf.gfile.ListDirectory(os.path.join(export_dir, "assets"))
self.assertListEqual(["tokens.txt"], module_files)
def make_encoder(params, is_training):
network_type = params['network_type']
if network_type == 'fully_connected':
encoder_ = fully_connected_encoder(params, is_training)
elif network_type == 'conv':
encoder_ = conv_encoder(params, is_training)
elif network_type == 'infoGAN':
encoder_ = infoGAN_encoder(params, is_training)
elif network_type == 'resnet':
encoder_ = resnet_encoder(params, is_training)
else:
raise NotImplementedError("Network type not implemented.")
def encoder_spec():
x = tf.compat.v1.placeholder(tf.float32, shape=params['full_size'])
z = encoder_(x)
hub.add_signature(inputs={'x': x}, outputs={'z': z})
enc_spec = hub.create_module_spec(encoder_spec)
encoder = hub.Module(enc_spec, name='encoder', trainable=True)
hub.register_module_for_export(encoder, "encoder")
return encoder
def export_hub_from_seved_model(model_export_dir, hub_dir):
# get module function and variable array dict
module_fn, variable_dict = get_mnist_module_fn(model_export_dir)
spec = hub.create_module_spec(module_fn=module_fn,
drop_collections=['losses', 'train_op'])
with tf.Graph().as_default():
m = hub.Module(spec)
print("------ check variables ------")
print(tf.global_variables())
print(variable_dict)
# assign variable to graph
init_w = tf.assign(m.variable_map["w"],
tf.convert_to_tensor(variable_dict["w/read:0"]))
init_b = tf.assign(m.variable_map["b"],
tf.convert_to_tensor(variable_dict["b/read:0"]))
init_vars = tf.group([init_w, init_b])
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
sess.run([init_vars])
# export TensorFlow hub module
m.export(hub_dir, sess)
def make_module_spec_for_testing(input_image_height=289,
input_image_width=289,
output_feature_dim=64):
"""Makes a stub image feature module for use in `TFImageProcessor` tests.
The resulting module has the signature expected by `TFImageProcessor`, but it
has no trainable variables and its initialization loads nothing from disk.
Args:
input_image_height: int, height of the module's input images.
input_image_width: int, width of module's input images.
output_feature_dim: int, dimension of the output feature vectors.
Returns:
`hub.ModuleSpec`
"""
def module_fn():
"""Builds the graph and signature for the stub TF-hub module."""
image_data = tf.placeholder(
shape=[1, input_image_height, input_image_width, 3], dtype=tf.float32)
# Linearly project image_data to shape [1, output_feature_dim] features.
projection_matrix = tf.ones([tf.size(image_data), output_feature_dim],
dtype=tf.float32)
encoder_output = tf.matmul(
tf.reshape(image_data, [1, -1]), projection_matrix)
# NB: the input feature must be named 'images' to satisfy
# hub.image_util.get_expected_image_size().
hub.add_signature(
'default', inputs={'images': image_data}, outputs=encoder_output)
return hub.create_module_spec(module_fn)
def make_decoder(params, is_training):
network_type = params['network_type']
if network_type == 'fully_connected':
decoder_ = fully_connected_decoder(params, is_training)
elif network_type == 'conv':
decoder_ = conv_decoder(params, is_training)
elif network_type == 'infoGAN':
decoder_ = infoGAN_decoder(params, is_training)
elif network_type == 'resnet':
decoder_ = resnet_decoder(params, is_training)
else:
raise NotImplementedError("Network type not implemented.")
def decoder_spec():
z = tf.compat.v1.placeholder(tf.float32,
shape=[None, params['latent_size']])
x = decoder_(z)
hub.add_signature(inputs={'z': z}, outputs={'x': x})
dec_spec = hub.create_module_spec(decoder_spec)
decoder = hub.Module(dec_spec, name='decoder', trainable=True)
hub.register_module_for_export(decoder, "decoder")
return decoder
def testMultipleApplicationsInDifferentScopes(self):
export_path = os.path.join(self.get_temp_dir(), "module-applied-in-scope")
spec = hub.create_module_spec(another_stateful_module_fn)
stateful_module = hub.Module(spec, name="moduleA")
with tf.name_scope("foo"):
with tf.variable_scope("bar"):
times2 = stateful_module(tf.constant([2.0]))
with tf.name_scope("baz"):
times3 = stateful_module(tf.constant([3.0]))
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
self.assertAllClose(sess.run(times2), [6.0])
self.assertAllClose(sess.run(times3), [9.0])
self.assertEqual(len(stateful_module.variable_map), 1)
self.assertEquals(
stateful_module.variable_map["iamtheoneandonly"].name,
"moduleA/iamtheoneandonly:0")
stateful_module.export(export_path, sess)
# Check minimal functionality of the exported module.
tf.reset_default_graph()
stateful_module = hub.Module(export_path, name="moduleB")
times2 = stateful_module(tf.constant([2.0]))
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
self.assertAllClose(sess.run(times2), [6.0])
def testApplyStatefulModuleMultipleTimes(self):
export_path = os.path.join(self.get_temp_dir(), "another-module")
with tf.Session() as sess:
spec = hub.create_module_spec(another_stateful_module_fn)
stateful_module = hub.Module(spec, trainable=True)
times2 = stateful_module(tf.constant([2.0]))
times3 = stateful_module(tf.constant([3.0]))
step = tf.Variable(0, trainable=False, name="global_step")
# Training will adapt the hidden variable to be approximately 2:
train = tf.contrib.layers.optimize_loss(
loss=tf.losses.mean_squared_error(times2, [4.0]),
global_step=step,
learning_rate=0.05,
optimizer="SGD")
sess.run(tf.global_variables_initializer())
for _ in range(50):
sess.run(train)
self.assertAllClose(sess.run(times2), [4.0])
self.assertAllClose(sess.run(times3), [6.0])
stateful_module.export(export_path, sess)
with tf.Session() as sess:
stateful_module = hub.Module(export_path)
times4 = stateful_module(tf.constant([4.0]))
times5 = stateful_module(tf.constant([5.0]))
sess.run(tf.global_variables_initializer())
self.assertAllClose(sess.run(times4), [8.0])
self.assertAllClose(sess.run(times5), [10.0])
def testDuplicateAssetCopy(self):
export_path = os.path.join(self.get_temp_dir(), "assets-module")
def module_with_duplicate_asset():
vocabulary_file = self.create_vocab_file("tokens2.txt", ["1", "2", "3"])
indices1 = tf.placeholder(dtype=tf.int64, name="indices1")
indices2 = tf.placeholder(dtype=tf.int64, name="indices2")
hub.add_signature(
inputs={
"indices_1": indices1,
"indices_2": indices2,
},
outputs={
"x": do_table_lookup(indices1, vocabulary_file),
"y": do_table_lookup(indices2, vocabulary_file),
})
with tf.Graph().as_default():
spec = hub.create_module_spec(module_with_duplicate_asset)
module_a = hub.Module(spec)
module_a({"indices_1": tf.constant([1, 2], dtype=tf.int64),
"indices_2": tf.constant([1, 2], dtype=tf.int64)}, as_dict=True)
with tf.Session() as sess:
sess.run(tf.tables_initializer())
module_a.export(export_path, sess)
def testLoadModuleFromFuncDef(self):
with tf.Session() as sess:
v = tf.placeholder(tf.int64)
spec = hub.create_module_spec(stateless_module_fn)
m = hub.Module(spec)
y = m(v)
self.assertEqual(sess.run(y, feed_dict={v: 10}), 100)
def testUpdateOps(self):
spec = hub.create_module_spec(update_ops_module_fn)
with tf.Session() as sess:
trainable_module = hub.Module(spec, trainable=True)
fixed_module = hub.Module(spec, trainable=False)
# TODO(b/62433105): Understand what is the desired behaviour of UPDATE_OPS
# and applying a Module multiple times. For now UPDATE_OPS probably only
# do something reasonable if each Module is applied exactly one time.
trainable_module()
fixed_module()
variable = tf.Variable(0.0)
step = tf.Variable(0, trainable=False, name="global_step")
train_op = tf.contrib.layers.optimize_loss(
variable,
global_step=step,
learning_rate=0.1,
optimizer="SGD")
sess.run(tf.global_variables_initializer())
sess.run(train_op)
trainable_module_vars = list(trainable_module.variable_map.values())
self.assertEqual(len(trainable_module_vars), 1)
self.assertEqual(sess.run(trainable_module_vars[0]), 1)
fixed_module_vars = list(fixed_module.variable_map.values())
self.assertEqual(len(fixed_module_vars), 1)
self.assertEqual(sess.run(fixed_module_vars[0]), 0)
def main(argv):
del argv # unused
tf.gfile.MakeDirs(FLAGS.export_dir)
def generative_model_fn():
code = hub.Module(FLAGS.flow_module)
decoder = hub.Module(FLAGS.decoder_module)
input_info = code.get_input_info_dict()
inputs = {
k: tf.placeholder(tf.float32, shape=input_info[k].get_shape())
for k in input_info.keys()
}
hub.add_signature(inputs=inputs, outputs=decoder(code(inputs)))
hub.attach_message(
"stamp_size",
decoder.get_attached_message("stamp_size", tf.train.Int64List))
hub.attach_message(
"pixel_size",
decoder.get_attached_message("pixel_size", tf.train.FloatList))
generative_model_spec = hub.create_module_spec(generative_model_fn)
generative_model = hub.Module(generative_model_spec, name="flow_module")
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
generative_model.export(FLAGS.export_dir + '/generator', sess)
print("Model exported in " + FLAGS.export_dir + '/generator')
def test_module_export_vocab_on_custom_fs(self):
root_dir = "file://%s" % self.get_temp_dir()
export_dir = "%s_%s" % (root_dir, "export")
tf_v1.gfile.MakeDirs(export_dir)
# Create a module with a vocab file located on a custom filesystem.
vocab_dir = os.path.join(root_dir, "vocab_location")
tf_v1.gfile.MakeDirs(vocab_dir)
vocab_filename = os.path.join(vocab_dir, "tokens.txt")
tf_utils.atomic_write_string_to_file(vocab_filename, "one", False)
def create_assets_module_fn():
def assets_module_fn():
indices = tf_v1.placeholder(dtype=tf.int64, name="indices")
table = index_to_string_table_from_file(
vocabulary_file=vocab_filename, default_value="UNKNOWN")
outputs = table.lookup(indices)
hub.add_signature(inputs=indices, outputs=outputs)
return assets_module_fn
with tf.Graph().as_default():
assets_module_fn = create_assets_module_fn()
spec = hub.create_module_spec(assets_module_fn)
embedding_module = hub.Module(spec)
with tf_v1.Session() as sess:
sess.run(tf_v1.tables_initializer())
embedding_module.export(export_dir, sess)
module_files = tf_v1.gfile.ListDirectory(export_dir)
self.assertListEqual(
["assets", "saved_model.pb", "tfhub_module.pb", "variables"],
sorted(module_files))
module_files = tf_v1.gfile.ListDirectory(os.path.join(export_dir, "assets"))
self.assertListEqual(["tokens.txt"], module_files)
def export_as_tfhub_module(hparams, problem, ckpt_dir, export_dir):
"""Exports the last checkpoint from the directory as tfhub module.
It creates the Module spec and signature (based on T2T problem information),
which is later used to create and export the hub module.
Module will be saved inside the ckpt_dir.
Args:
hparams: T2T parameters, model graph will be based on them.
problem: the name of the problem
ckpt_dir: directory with the checkpoints.
export_dir: Directory to write the exported model to.
"""
def hub_module_fn():
"""Creates the TF graph for the hub module."""
model_fn = t2t_model.T2TModel.make_estimator_model_fn(
FLAGS.model,
hparams,
decode_hparams=decoding.decode_hparams(FLAGS.decode_hparams))
features = problem.serving_input_fn(hparams).features
spec = model_fn(features, labels=None, mode=tf.estimator.ModeKeys.PREDICT)
# Currently only supports a single input and single output.
hub.add_signature(
inputs=features, outputs=spec.export_outputs["serving_default"].outputs)
module_spec = hub.create_module_spec(hub_module_fn)
# Loads the weights from the checkpoint using the model above
# and saves it in the export_path.
export_module_spec_with_checkpoint(
module_spec,
checkpoint_path=tf.train.latest_checkpoint(ckpt_dir),
export_path=export_dir,
scope_prefix="")
def export_functional_and_derivatives(
self,
export_path: str,
batch_dim: Optional[int] = None,
):
"""Exports the TensorFlow graph containing the functional and derivatives.
The hub modules supplied contain the TensorFlow operations for the
evaluation of the exchange-correlation energy. Evaluation of the functional
derivatives, required for a self-consistent calculation, are added in
_build_graph. The module created by export_functional_and_derivatives
contains the evaluation of the functional and the functional derivatives.
This is much simpler to use from languages other than Python, e.g. using the
C or C++ TensorFlow API, or using tfcompile to create a standalone C++
library.
Args:
export_path: path to write the Hub model to. The exported model can be
loaded using either TF-Hub or SavedModel APIs.
batch_dim: the batch dimension of the grid to use in the model. Default:
None (determine at runtime). This should only be set if the exported
model is to be ahead-of-time compiled into a standalone library.
"""
with tf.Graph().as_default():
spec = hub.create_module_spec(
self._build_graph, tags_and_args=[(set(), {'batch_dim': batch_dim})])
functional_and_derivatives = hub.Module(spec=spec)
with tf.Session() as session:
session.run(tf.global_variables_initializer())
functional_and_derivatives.export(export_path, session)
def main():
load_shufflenet()
spec = hub.create_module_spec(module_fn)
module = hub.Module(spec)
with tf.Session() as sess:
module.export("onnx/shufflenet/1", sess)
print("Exported")
def testModuleWithRegularizedLayers(self):
# The linear map y = Mx + b with L2 regularization on M and b
# when trained at x = [1,1] with L2 loss towards the target y' = [4,4]
# learns M = [[1,1],[1,1]], b = [1,1], y = [3,3], with eight balanced
# loss terms: the elements of M, b, and y' - y are all distance 1 from zero.
train_input = [[1.0, 1.0]]
target = [[4.0, 4.0]]
spec = hub.create_module_spec(layers_module_fn)
with tf.Graph().as_default():
m = hub.Module(spec, trainable=True)
x = tf.placeholder(dtype=tf.float32)
y = m(x)
squared_loss = tf.losses.mean_squared_error(y, target, weights=2.0)
# Recover REGULARIZATION_LOSSES from the module.
total_loss = squared_loss + tf.losses.get_regularization_loss()
step = tf.Variable(0, trainable=False, name="global_step")
train = tf.contrib.layers.optimize_loss(
loss=total_loss, global_step=step, learning_rate=0.1, optimizer="SGD")
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for _ in range(50):
sess.run(train, feed_dict={x: train_input})
# Verify M = [[1,1],[1,1]], b = [1,1] by evaluating at three points.
# Without regularization, the result would be an underdetermined mess.
out = sess.run(y, feed_dict={x: [[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]]})
self.assertAllClose(
out, [[1.0, 1.0], [2.0, 2.0], [2.0, 2.0]], atol=0.001)
def model_fn(features, labels, mode, params):
training = mode == tf.estimator.ModeKeys.TRAIN
spec = tfhub.create_module_spec(module_fn,
tags_and_args=[
({'train'}, {
'training': True,
'params': params
}),
(set(), {
'training': False,
'params': params
}),
])
tags = {'train'} if training else None
module = tfhub.Module(spec, trainable=training, tags=tags)
tfhub.register_module_for_export(module, 'doodle')
image = features['image']
image.shape.assert_is_compatible_with([None, 28, 28, 1])
predictions = module(image, as_dict=True)
if mode == tf.estimator.ModeKeys.PREDICT:
return tf.estimator.EstimatorSpec(
mode=mode,
predictions=predictions,
export_outputs={
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
tf.estimator.export.PredictOutput(predictions),
})
with tf.variable_scope('losses'):
cross_entropy_loss = tf.losses.sparse_softmax_cross_entropy(
labels=labels, logits=predictions['logits'])
total_loss = tf.losses.get_total_loss()
tf.summary.image('image', image)
tf.summary.scalar('total_loss', total_loss)
metric_ops = metrics.calculate(labels, predictions['classes'],
params['num_classes'])
if mode == tf.estimator.ModeKeys.EVAL:
return tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss,
eval_metric_ops=metric_ops)
if mode == tf.estimator.ModeKeys.TRAIN:
global_step = tf.train.get_or_create_global_step()
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.variable_scope('optimizer'):
optimizer = tf.train.AdamOptimizer(params['learning_rate'])
with tf.control_dependencies(update_ops):
fit = optimizer.minimize(total_loss, global_step)
return tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss,
train_op=fit,
eval_metric_ops=metric_ops)
def testUnusedInputModule(self):
with tf.Session() as sess:
v1 = tf.placeholder(tf.int64)
v2 = tf.placeholder(tf.int64)
spec = hub.create_module_spec(unused_input_module_fn)
m = hub.Module(spec)
out = m({"x": v1, "unused": v2})
self.assertEqual(sess.run(out, feed_dict={v1: 10, v2: 4}), 100)
def build_hub_module(model, num_classes, global_step, checkpoint_path):
"""Create TF-Hub module."""
tags_and_args = [
# The default graph is built with batch_norm, dropout etc. in inference
# mode. This graph version is good for inference, not training.
([], {'is_training': False}),
# A separate "train" graph builds batch_norm, dropout etc. in training
# mode.
(['train'], {'is_training': True}),
]
def module_fn(is_training):
"""Function that builds TF-Hub module."""
endpoints = {}
inputs = tf.placeholder(
tf.float32, [None, None, None, 3])
with tf.variable_scope('base_model', reuse=tf.AUTO_REUSE):
hiddens = model(inputs, is_training)
for v in ['initial_conv', 'initial_max_pool', 'block_group1',
'block_group2', 'block_group3', 'block_group4',
'final_avg_pool']:
endpoints[v] = tf.get_default_graph().get_tensor_by_name(
'base_model/{}:0'.format(v))
if FLAGS.train_mode == 'pretrain':
hiddens_proj = model_util.projection_head(hiddens, is_training)
endpoints['proj_head_input'] = hiddens
endpoints['proj_head_output'] = hiddens_proj
else:
logits_sup = model_util.supervised_head(
hiddens, num_classes, is_training)
endpoints['logits_sup'] = logits_sup
hub.add_signature(inputs=dict(images=inputs),
outputs=dict(endpoints, default=hiddens))
# Drop the non-supported non-standard graph collection.
drop_collections = ['trainable_variables_inblock_%d'%d for d in range(6)]
spec = hub.create_module_spec(module_fn, tags_and_args, drop_collections)
hub_export_dir = os.path.join(FLAGS.model_dir, 'hub')
checkpoint_export_dir = os.path.join(hub_export_dir, str(global_step))
if tf.io.gfile.exists(checkpoint_export_dir):
# Do not save if checkpoint already saved.
tf.io.gfile.rmtree(checkpoint_export_dir)
spec.export(
checkpoint_export_dir,
checkpoint_path=checkpoint_path,
name_transform_fn=None)
if FLAGS.keep_hub_module_max > 0:
# Delete old exported Hub modules.
exported_steps = []
for subdir in tf.io.gfile.listdir(hub_export_dir):
if not subdir.isdigit():
continue
exported_steps.append(int(subdir))
exported_steps.sort()
for step_to_delete in exported_steps[:-FLAGS.keep_hub_module_max]:
tf.io.gfile.rmtree(os.path.join(hub_export_dir, str(step_to_delete)))
def testVariables(self):
spec = hub.create_module_spec(stateful_module_fn)
m = hub.Module(spec, name="test")
out = m()
self.assertEqual(list(m.variable_map.keys()), ["var123"])
self.assertEqual(m.variable_map["var123"].name, "test/var123:0")
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
self.assertAllClose(sess.run(out), [1.0, 2.0, 3.0])
def flow_model_fn(features, labels, mode, params, config):
"""
Model function to create a VAE estimator
"""
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
if mode == tf.estimator.ModeKeys.PREDICT:
y = features
def flow_module_spec():
cond_layer = tf.placeholder(tf.float32, shape=[None, n_cond])
flow = params['flow_fn'](cond_layer, is_training)
hub.add_signature(inputs=cond_layer,
outputs=flow.sample(tf.shape(cond_layer)[0]))
flow_spec = hub.create_module_spec(flow_module_spec)
flow = hub.Module(flow_spec, name='flow_module')
hub.register_module_for_export(flow, "code_sampler")
predictions = {'code': flow(y)}
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
x = features['x']
y = features['y']
# Loads the encoding function to work on the images
code = x
with tf.variable_scope("flow_module"):
cond_layer = y
flow = params['flow_fn'](cond_layer, is_training)
loglikelihood = flow.log_prob(code)
# This is the loglikelihood of a batch of images
tf.summary.scalar('loglikelihood', tf.reduce_mean(loglikelihood))
loss = -tf.reduce_mean(loglikelihood)
# Training of the model
global_step = tf.train.get_or_create_global_step()
learning_rate = tf.train.cosine_decay(params["learning_rate"], global_step,
params["max_steps"])
tf.summary.scalar("learning_rate", learning_rate)
optimizer = tf.train.AdamOptimizer(learning_rate)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(loss, global_step=global_step)
eval_metric_ops = {
"loglikelihood": tf.metrics.mean(tf.reduce_mean(loglikelihood))
}
return tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op,
eval_metric_ops=eval_metric_ops)
def body(self, features):
hparams = self.hparams
model_opt = {
'nr_resnet': hparams.nr_resnet,
'nr_filters': hparams.hidden_size,
'nr_logistic_mix': 1,
'resnet_nonlinearity': 'concat_elu',
'energy_distance': False,
'dropout_p': hparams.dropout
}
def pixel_cnn_fn(input_layer):
model = tf.make_template('model', model_spec)
out = model(input_layer, None, ema=None, **model_opt)
out = tf.layers.dense(out, 2, activation=None)
loc, scale = tf.split(out, num_or_size_splits=2, axis=-1)
scale = tf.nn.softplus(scale) + 1e-4
distribution = tfp.distributions.Independent(
tfp.distributions.Normal(loc=loc, scale=scale))
sample = distribution.sample()
log_prob = distribution.log_prob(input_layer)
grads = tf.gradients(log_prob, input_layer)[0]
print(grads)
output = {
'sample': sample,
'log_prob': log_prob,
'logits': out,
'mu': loc,
'scale': scale,
'grads': grads
}
return output
# During predict, we use a tf hub module, otherwise we stick to normal
# behavior to allow for multi gpu training
if hparams.mode == tf.estimator.ModeKeys.PREDICT:
def make_model_spec():
input_layer = tf.placeholder(
tf.float32, shape=features["inputs"].get_shape())
outputs = pixel_cnn_fn(input_layer)
hub.add_signature(inputs=input_layer, outputs=outputs)
spec = hub.create_module_spec(make_model_spec,
drop_collections=['checkpoints'])
pixelcnn = hub.Module(spec, name="pixelcnn_module", trainable=True)
hub.register_module_for_export(pixelcnn, "pixelcnn")
output = pixelcnn(features["inputs"], as_dict=True)
else:
with tf.variable_scope('pixelcnn_module'):
output = pixel_cnn_fn(features["inputs"])
self.image_summary("inputs", features["inputs"])
self.image_summary("samples", output["sample"])
return output["logits"], {"training": -output["log_prob"]}
def exportHubModule(self, path, hub_model_fn):
import tensorflow_hub as hub
tf.reset_default_graph()
spec = hub.create_module_spec(hub_model_fn)
module = hub.Module(spec, name=self.name, trainable=True)
with self.makeSession() as sess:
module.export(path, sess)
def testMultipleOutputs(self):
with tf.Session() as sess:
spec = hub.create_module_spec(multiple_outputs_module_fn)
m = hub.Module(spec)
output = m(tf.constant([2.0]), as_dict=True)
output1 = output["y"]
output2 = output["z"]
sess.run(tf.global_variables_initializer())
self.assertAllClose(sess.run(output1), [6.0])
self.assertAllClose(sess.run(output2), [18.0])
def testConvertToTensor(self):
spec = hub.create_module_spec(stateless_module_fn)
with tf.Session() as sess:
m = hub.Module(spec)
y = m([10, 2])
self.assertAllEqual(sess.run(y), [100, 4])
with tf.Session() as sess:
m = hub.Module(spec)
with self.assertRaises(TypeError):
m("hello")
def _save_plus_one_hub_module_v1(path):
"""Writes TF1.x hub.Module that increments the input by one."""
def plus_one():
x = tf.compat.v1.placeholder(dtype=tf.float32, name="x")
y = x + 1
hub.add_signature(inputs=x, outputs=y)
spec = hub.create_module_spec(plus_one)
_export_module_spec_with_init_weights(spec, path)
def testImageSizeManuallySpecified(self):
spec = hub.create_module_spec(create_image_module_fn([None, None]))
image_column = hub.image_embedding_column("image",
spec,
image_size=[229, 229])
parsing_spec = tf_v1.feature_column.make_parse_example_spec(
[image_column])
self.assertEqual(
parsing_spec,
{"image": tf_v1.FixedLenFeature([229, 229, 3], dtype=tf.float32)})
def _plus_one_model_tf1():
def plus_one():
x = tf.compat.v1.placeholder(dtype=tf.float32, name="x")
y = x + 1
hub.add_signature(inputs=x, outputs=y)
spec = hub.create_module_spec(plus_one)
with tf.compat.v1.get_default_graph().as_default():
module = hub.Module(spec, trainable=True)
return module
def testLargePartitionedVariables(self):
spec = hub.create_module_spec(
create_partitioned_variable_module_fn(partitions=25, shape=[600, 3]))
m = hub.Module(spec, name="test")
out = m()
self.assertEqual(len(m.variable_map), 2)
self.assertEqual(len(m.variable_map["partitioned_variable"]), 25)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
self.assertAllClose(sess.run(out), 2 * np.ones([600, 3]))
def main(_):
tags_and_args = []
for is_training in (True, False):
tags = set()
if is_training:
tags.add("train")
tags_and_args.append((tags, dict(is_training=is_training)))
spec = hub.create_module_spec(module_fn, tags_and_args=tags_and_args)
checkpoint_path = os.path.join(FLAGS.albert_directory, FLAGS.checkpoint_name)
tf.logging.info("Using checkpoint {}".format(checkpoint_path))
spec.export(FLAGS.export_path, checkpoint_path=checkpoint_path)
def build_hub_module(yml_config, num_classes, hub_id_name, checkpoint_path,
save_path):
"""Create TF-Hub module."""
tags_and_args = [
# The default graph is built with batch_norm, dropout etc. in inference
# mode. This graph version is good for inference, not training.
([], {
'is_training': False
}),
# A separate "train" graph builds batch_norm, dropout etc. in training
# mode.
(['train'], {
'is_training': True
}),
]
def module_fn(is_training):
endpoints = {}
inputs = tf1.placeholder(tf1.float32, [None, None, None, 3])
# Load the base network and set it to non-trainable (for speedup fine-tuning)
hub_path = str(
Path(yml_config['finetuning']['pretrained_build']).resolve())
hub_path = os.path.join(hub_path, 'hub')
module = hub.Module(hub_path, trainable=is_training)
if yml_config['finetuning']['pretrained_model'] == 'ChestXRay':
key = module(inputs=inputs,
signature="projection-head-1",
as_dict=True)
else:
key = module(inputs=inputs, as_dict=True)
# Attach a trainable linear layer to adapt for the new task.
with tf1.variable_scope('head_supervised_new', reuse=tf1.AUTO_REUSE):
logits_t = tf1.layers.dense(inputs=key['default'],
units=num_classes)
endpoints['head_classification'] = logits_t
hub.add_signature(inputs=dict(images=inputs),
outputs=dict(endpoints, default=logits_t))
spec = hub.create_module_spec(module_fn, tags_and_args)
hub_path = str(Path(save_path) / 'hub')
if os.path.exists(hub_path):
rmtree(hub_path)
checkpoint_export_dir = os.path.join(save_path, 'hub')
create_folder(checkpoint_export_dir)
spec.export(checkpoint_export_dir,
checkpoint_path=checkpoint_path,
name_transform_fn=None)
return checkpoint_export_dir
def testModuleWithMultipleSignatures(self):
spec = hub.create_module_spec(multiple_signature_module_fn)
module_a = hub.Module(spec, name="moduleA")
in_tensor = tf.placeholder(dtype=tf.float32)
out_tensor_a = module_a(in_tensor, signature="mul")
out_tensor_b = module_a(out_tensor_a, signature="div")
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
in_values = [6, 3, 1]
self.assertAllClose(
sess.run(out_tensor_b, feed_dict={in_tensor: in_values}), in_values)
def exportHubModule(self, path, input_shape, trainable=False):
self.input_shape = input_shape
spec = hub.create_module_spec(self.hub_model_fn)
module = hub.Module(spec, name=Estimator.NAME, trainable=trainable)
with tf.Session() as sess:
saver = tf.train.Saver()
saver.restore(sess, self.estimator.latest_checkpoint())
sess.run(tf.tables_initializer())
module.export(path, sess)
def _make_module_spec(keys_vocab, embedding_size, is_trainable, combiner,
max_norm, shard_bytes):
def module_fn():
embeddings_shape = [len(keys_vocab), embedding_size]
embedding_weights = tf.get_variable(
name=_EMBEDDINGS_VAR_NAME,
shape=embeddings_shape,
dtype=tf.float32,
initializer=tf.zeros_initializer(),
trainable=is_trainable,
partitioner=tf.variable_axis_size_partitioner(
max_shard_bytes=shard_bytes # 1Gb by default
))
lookup_table = tf.contrib.lookup.index_table_from_tensor(
mapping=keys_vocab, default_value=0)
default_keys = tf.placeholder(dtype=tf.string, shape=[None])
default_ids = lookup_table.lookup(default_keys)
default_embeddings = tf.nn.embedding_lookup(
params=embedding_weights,
ids=default_ids,
partition_strategy='div',
max_norm=max_norm,
)
hub.add_signature('default', default_keys, default_embeddings)
context_keys = tf.sparse_placeholder(dtype=tf.string,
shape=(None, None))
context_ids = lookup_table.lookup(context_keys)
context_embeddings = tf.nn.safe_embedding_lookup_sparse(
embedding_weights=embedding_weights,
sparse_ids=context_ids,
combiner=combiner,
default_id=0,
partition_strategy='div',
max_norm=max_norm,
)
hub.add_signature('context', context_keys, context_embeddings)
sequence_keys = tf.sparse_placeholder(dtype=tf.string,
shape=(None, None, None))
sequence_ids = lookup_table.lookup(sequence_keys)
sequence_embeddings = tf.nn.safe_embedding_lookup_sparse(
embedding_weights=embedding_weights,
sparse_ids=sequence_ids,
combiner=combiner,
default_id=0,
partition_strategy='div',
max_norm=max_norm,
)
hub.add_signature('sequence', sequence_keys, sequence_embeddings)
return hub.create_module_spec(module_fn)
def testWhileModule(self):
spec = hub.create_module_spec(while_module_fn)
with tf.Graph().as_default():
with tf.Session() as sess:
x = tf.placeholder(tf.float32)
n = tf.placeholder(tf.int32)
pow_module = hub.Module(spec)
y = pow_module({"x": x, "n": n})
self.assertAllClose(sess.run(y, {x: 9.1, n: 1}), 9.1)
self.assertAllClose(sess.run(y, {x: 2.4, n: 2}), 5.76)
grad = tf.gradients([y], [x])
self.assertAllClose(sess.run(grad, {x: 2, n: 3}), [12.0])
def _generate_module(self):
spec = hub.create_module_spec(self._stateless_module_fn)
m = hub.Module(spec, name="test_module")
out = m(10)
export_path = os.path.join(self.get_temp_dir(), "module")
with tf_v1.Session() as sess:
sess.run(tf_v1.global_variables_initializer())
self.assertAllClose(sess.run(out), 100)
m.export(export_path, sess)
self._create_tgz(export_path)
def _testReluModule(self, module_fn):
spec = hub.create_module_spec(module_fn)
with tf.Graph().as_default():
with tf.Session() as sess:
x = tf.placeholder(dtype=tf.float32, name="x")
relu_module = hub.Module(spec)
y = relu_module(x)
self.assertAllClose(sess.run(y, {x: 9.1}), 9.1)
self.assertAllClose(sess.run(y, {x: -2.4}), 0.0)
grad = tf.gradients([y], [x])
self.assertAllClose(sess.run(grad, {x: 2}), [1.0])
self.assertAllClose(sess.run(grad, {x: -2}), [0.0])
def ExportFakeTF1ImageModule(*, input_image_height: int,
input_image_width: int, output_feature_dim: int,
export_path: str):
"""Makes a TF-hub image feature module for use in unit tests.
The resulting module has the signature of a image model, but contains a
minimal set of trainable variables and its initialization loads nothing from
disk.
Args:
input_image_height: Height of the module's input images.
input_image_width: Width of module's input images.
output_feature_dim: Dimension of the output feature vectors.
export_path: Path where exported module will be written.
"""
def ModuleFn(training):
"""Builds the graph and signature for the stub TF-hub module."""
image_data = tf.placeholder(
shape=[None, input_image_height, input_image_width, 3],
dtype=tf.float32)
# Linearly project image_data to shape [1, output_feature_dim] features.
encoder_output = tf.compat.v1.layers.dense(
tf.reshape(image_data,
[-1, input_image_height * input_image_width * 3]),
output_feature_dim)
# Add a non-trainable 'count' variable that can be updated through an
# UPDATE_OP. This is analogous to a batch-norm moving average that should be
# updated during fine-tuning.
v = tf.get_variable('count',
initializer=0,
dtype=tf.int32,
trainable=False)
if training:
update_op = v.assign_add(1).op
tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_op)
hub.add_signature('default',
inputs={'images': image_data},
outputs=encoder_output)
spec = hub.create_module_spec(ModuleFn,
tags_and_args=[
({'train'}, dict(training=True)),
(set(), dict(training=False))
])
with tf.compat.v1.Graph().as_default():
module = hub.Module(spec, trainable=True)
with tf.compat.v1.Session() as session:
session.run(tf.compat.v1.global_variables_initializer())
module.export(export_path, session)
def main(_):
tf.logging.info("Running export_to_tfhub.py")
tags_and_args = []
for is_training in (True, False):
tags = set()
if is_training:
tags.add("train")
tags_and_args.append((tags, dict(is_training=is_training)))
spec = hub.create_module_spec(module_fn, tags_and_args=tags_and_args)
tf.logging.info("Using checkpoint {}".format(FLAGS.checkpoint_path))
tf.logging.info("Exporting to {}".format(FLAGS.export_path))
spec.export(FLAGS.export_path, checkpoint_path=FLAGS.checkpoint_path)
def create_module():
tf.compat.v1.logging.set_verbosity(tf.logging.INFO)
tags_and_args = []
tags = set()
tags_and_args.append((tags, dict(is_training=False)))
module_fn = build_module_fn(BERT_CONFIG, BERT_VOCAB, LOWER_CASE)
spec = hub.create_module_spec(module_fn, tags_and_args=tags_and_args)
try:
spec.export(MODULE_FOLDER,
checkpoint_path=BERT_INIT_CHKPNT)
except tf.errors.AlreadyExistsError:
tf.compat.v1.logging.warning("Path already exists")
def export_module(module_export_path):
"""Create and export a simple module to the specified path."""
def _stateless_module_fn():
"""Simple module that squares an input."""
x = tf.compat.v1.placeholder(tf.int64)
y = x * x
hub.add_signature(inputs=x, outputs=y)
spec = hub.create_module_spec(_stateless_module_fn)
m = hub.Module(spec, name="test_module")
with tf.compat.v1.Session() as sess:
sess.run(tf.compat.v1.global_variables_initializer())
m.export(module_export_path, sess)
def export_module(path):
spec = hub.create_module_spec(half_plus_two)
with tf.Graph().as_default():
module = hub.Module(spec)
init_a = tf.assign(module.variable_map["a"], 0.5)
init_b = tf.assign(module.variable_map["b"], 2.0)
init_vars = tf.group([init_a, init_b])
with tf.Session() as session:
session.run(init_vars)
module.export(path, session)
def export_module(path):
spec = hub.create_module_spec(half_plus_two)
with tf.Graph().as_default():
module = hub.Module(spec)
init_a = tf.assign(module.variable_map["a"], 0.5)
init_b = tf.assign(module.variable_map["b"], 2.0)
init_vars = tf.group([init_a, init_b])
with tf.Session() as session:
session.run(init_vars)
module.export(path, session)
def _save_plus_one_hub_module_v1(path):
def plus_one():
x = tf.compat.v1.placeholder(dtype=tf.float32, name='x')
y = x + 1
hub.add_signature(inputs=x, outputs=y)
spec = hub.create_module_spec(plus_one)
with tf.compat.v1.Graph().as_default():
module = hub.Module(spec, trainable=True)
with tf.compat.v1.Session() as session:
session.run(tf.compat.v1.global_variables_initializer())
module.export(path, session)
def export_as_tfhub_module(model_name,
hparams,
decode_hparams,
problem,
checkpoint_path,
export_dir):
"""Exports the last checkpoint from the directory as tfhub module.
It creates the Module spec and signature (based on T2T problem information),
which is later used to create and export the hub module.
Module will be saved inside the ckpt_dir.
Args:
model_name: name of the model to be exported.
hparams: T2T parameters, model graph will be based on them.
decode_hparams: T2T parameters for decoding.
problem: the name of the problem
checkpoint_path: path to the checkpoint to be exported.
export_dir: Directory to write the exported model to.
"""
def hub_module_fn():
"""Creates the TF graph for the hub module."""
model_fn = t2t_model.T2TModel.make_estimator_model_fn(
model_name,
hparams,
decode_hparams=decode_hparams,
use_tpu=FLAGS.use_tpu)
features = problem.serving_input_fn(hparams).features
# we must do a copy of the features, as the model_fn can add additional
# entries there (like hyperparameter settings etc).
original_features = features.copy()
spec = model_fn(features, labels=None, mode=tf.estimator.ModeKeys.PREDICT)
hub.add_signature(
inputs=original_features,
outputs=spec.export_outputs["serving_default"].outputs)
# TFHub doesn't support the following collections.
drop_collections = [tf.GraphKeys.LOSSES,
tf.GraphKeys.SUMMARIES, tf.GraphKeys.LOCAL_VARIABLES]
module_spec = hub.create_module_spec(
hub_module_fn, drop_collections=drop_collections)
# Loads the weights from the checkpoint using the model above
# and saves it in the export_path.
export_module_spec_with_checkpoint(
module_spec,
checkpoint_path=checkpoint_path,
export_path=export_dir,
scope_prefix="")
def export_as_tfhub_module(model_name, hparams, decode_hparams, problem,
checkpoint_path, export_dir):
"""Exports the last checkpoint from the directory as tfhub module.
It creates the Module spec and signature (based on T2T problem information),
which is later used to create and export the hub module.
Module will be saved inside the ckpt_dir.
Args:
model_name: name of the model to be exported.
hparams: T2T parameters, model graph will be based on them.
decode_hparams: T2T parameters for decoding.
problem: the name of the problem
checkpoint_path: path to the checkpoint to be exported.
export_dir: Directory to write the exported model to.
"""
def hub_module_fn():
"""Creates the TF graph for the hub module."""
model_fn = t2t_model.T2TModel.make_estimator_model_fn(
model_name,
hparams,
decode_hparams=decode_hparams,
use_tpu=FLAGS.use_tpu)
features = problem.serving_input_fn(hparams,
decode_hparams,
use_tpu=FLAGS.use_tpu).features
# we must do a copy of the features, as the model_fn can add additional
# entries there (like hyperparameter settings etc).
original_features = features.copy()
spec = model_fn(features,
labels=None,
mode=tf_estimator.ModeKeys.PREDICT)
hub.add_signature(
inputs=original_features,
outputs=spec.export_outputs["serving_default"].outputs)
# TFHub doesn't support the following collections.
drop_collections = [
tf.GraphKeys.LOSSES, tf.GraphKeys.SUMMARIES,
tf.GraphKeys.LOCAL_VARIABLES
]
module_spec = hub.create_module_spec(hub_module_fn,
drop_collections=drop_collections)
# Loads the weights from the checkpoint using the model above
# and saves it in the export_path.
export_module_spec_with_checkpoint(module_spec,
checkpoint_path=checkpoint_path,
export_path=export_dir,
scope_prefix="")
def testVariableColocationPropagation(self):
spec = hub.create_module_spec(stateful_module_fn_with_colocation)
m = hub.Module(spec)
u1 = tf.constant(1, name="u1")
u2 = tf.constant(2, name="u2")
with tf.colocate_with(u1), tf.colocate_with(u2):
x = tf.constant(100.0, name="x")
y = m(x)
self.assertItemsEqual(y.op.colocation_groups(),
[tf.compat.as_bytes("loc:@module/var123"),
tf.compat.as_bytes("loc:@u1"),
tf.compat.as_bytes("loc:@u2")])
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
self.assertEqual(sess.run(y), 101.0)
def testExportedConsumerModelWorksIfItUsesHubModuleWithAssets(self):
# 1. Create and export a module with assets.
module_export_path = os.path.join(self.get_temp_dir(), "small-module")
vocabulary_file = self.create_vocab_file("tokens.txt",
["emerson", "lake", "palmer"])
assets_module_fn = create_assets_module_fn(vocabulary_file)
spec = hub.create_module_spec(assets_module_fn)
with tf.Graph().as_default():
small_module = hub.Module(spec)
with tf.Session() as sess:
small_module.export(module_export_path, sess)
# 2. Remove the original vocab file and move the module to another location.
tf.gfile.Remove(vocabulary_file)
inner_module_path = os.path.join(self.get_temp_dir(), "inner-module")
tf.gfile.Rename(module_export_path, inner_module_path)
del module_export_path
# 3. Use the module in a consumer model (which is another module here).
module_export_path = os.path.join(self.get_temp_dir(), "consumer-module")
consumer_module_fn = create_consumer_module_fn(inner_module_path)
spec = hub.create_module_spec(consumer_module_fn)
with tf.Graph().as_default():
consumer_module = hub.Module(spec)
with tf.Session() as sess:
consumer_module.export(module_export_path, sess)
# 4. Delete the inner module on disk and move the consumer model to a final
# location for serving.
tf.gfile.DeleteRecursively(inner_module_path)
module_serving_path = os.path.join(self.get_temp_dir(), "serving-module")
tf.gfile.Rename(module_export_path, module_serving_path)
# 5. Make sure the model can be served successfully.
with tf.Graph().as_default():
serving_module = hub.Module(module_serving_path)
output = serving_module(tf.constant([1, 2], dtype=tf.int64))
with tf.Session() as sess:
sess.run(tf.tables_initializer())
self.assertAllEqual(list(sess.run(output)), [b"lake", b"palmer"])
def testLoadTrainableModuleFromFuncDef(self):
with tf.Session() as sess:
spec = hub.create_module_spec(stateful_module_fn)
m = hub.Module(spec, trainable=True)
x = m()
step = tf.Variable(0, trainable=False, name="global_step")
train = tf.contrib.layers.optimize_loss(
loss=tf.losses.mean_squared_error(x, [3.1, 3.2, 3.3]),
global_step=step,
learning_rate=0.40,
optimizer="SGD")
sess.run(tf.global_variables_initializer())
for _ in range(50):
sess.run(train)
got = sess.run(x)
self.assertAllClose(got, [3.1, 3.2, 3.3])
def testPartitionedVariables(self):
spec = hub.create_module_spec(
create_partitioned_variable_module_fn(partitions=3, shape=[7, 3]))
m = hub.Module(spec, name="test")
out = m()
self.assertEqual(len(m.variable_map), 2)
self.assertEqual(m.variable_map["normal_variable"].name,
"test/normal_variable:0")
self.assertAllEqual(
[variable.name for variable in m.variable_map["partitioned_variable"]],
["test/partitioned_variable/part_0:0",
"test/partitioned_variable/part_1:0",
"test/partitioned_variable/part_2:0"])
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
self.assertAllClose(sess.run(out), 2 * np.ones([7, 3]))
def testNonResourceVariableInCond(self):
spec = hub.create_module_spec(stateful_non_rv_module_fn)
m = hub.Module(spec)
pred = tf.placeholder(tf.bool)
def true_fn():
v = m()
self.assertItemsEqual(v.op.colocation_groups(),
[tf.compat.as_bytes("loc:@module/var123")])
return v
def false_fn():
return tf.constant(9.0)
out = tf.cond(pred, true_fn, false_fn)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
self.assertEqual(sess.run(out, feed_dict={pred: True}), 10.0)
self.assertEqual(sess.run(out, feed_dict={pred: False}), 9.0)
def testNonResourceVariableInWhileLoop(self):
# This test uses non-Resource variables to see an actual colocation
# constraint propagated to the context Enter op. The long comment on
# colocation in testResourceVariables explains why they may not offer that.
spec = hub.create_module_spec(stateful_non_rv_module_fn)
m = hub.Module(spec)
cond = lambda i, x: tf.less(i, 4)
def body(i, x):
v = m()
self.assertItemsEqual(v.op.colocation_groups(),
[tf.compat.as_bytes("loc:@module/var123")])
return (tf.add(i, 1), 2*x)
oi, ox = tf.while_loop(cond, body, [0, 10.0])
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
self.assertAllEqual(sess.run([oi, ox]), [4, 160.0])
def testNestedControlFlowModule(self):
spec = hub.create_module_spec(nested_control_flow_module_fn)
with tf.Graph().as_default():
with tf.Session() as sess:
elems = tf.placeholder(tf.float32, shape=[None])
a = tf.placeholder(tf.float32)
m = hub.Module(spec)
out = m({"elems": elems, "a": a})
self.assertAllClose(
sess.run(out, {
a: 1.1,
elems: [10, 0, 0.5, 1.2]
}), 11.2)
grad = tf.gradients([out], [elems])
self.assertAllClose(sess.run(grad, {a: 1, elems: [10, 0, 0.5, 1.2]}),
[[1.0, 0.0, 0.0, 1.0]])
def testModuleWithTable(self):
export_path = os.path.join(self.get_temp_dir(), "table-module")
with tf.Graph().as_default():
spec = hub.create_module_spec(table_lookup_module_fn)
m = hub.Module(spec)
# Export requires a session to work regardless of the module having no
# variables to export.
with tf.Session() as sess:
m.export(export_path, sess)
with tf.Graph().as_default():
v = tf.placeholder(dtype=tf.int64)
f = hub.Module(export_path)
y = f(v)
with tf.Session() as sess:
sess.run(tf.tables_initializer())
got = sess.run(y, feed_dict={v: [0, 1, 2, 3]})
self.assertAllEqual(list(got), [b"index0", b"hello", b"world", b"UNK"])
def test_http_locations(self):
spec = hub.create_module_spec(self._stateless_module_fn)
m = hub.Module(spec, name="test_module")
out = m(10)
export_path = os.path.join(self.get_temp_dir(), "module")
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
self.assertAllClose(sess.run(out), 100)
m.export(export_path, sess)
os.chdir(export_path)
tar = tarfile.open("test_module.tgz", "w")
for f in self._list_module_files(export_path):
tar.add(f)
tar.close()
m = hub.Module("http://localhost:%d/test_module.tgz" % self.server_port)
out = m(11)
with tf.Session() as sess:
self.assertAllClose(sess.run(out), 121)
# Test caching using custom filesystem (file://) to make sure that the
# TF Hub library can operate on such paths.
try:
root_dir = "file://%s" % self.get_temp_dir()
cache_dir = "%s_%s" % (root_dir, "cache")
tf.gfile.MakeDirs(cache_dir)
os.environ["TFHUB_CACHE_DIR"] = cache_dir
m = hub.Module("http://localhost:%d/test_module.tgz" % self.server_port)
out = m(11)
with tf.train.MonitoredSession() as sess:
self.assertAllClose(sess.run(out), 121)
cache_content = sorted(tf.gfile.ListDirectory(cache_dir))
tf.logging.info("Cache context: %s", str(cache_content))
self.assertEqual(2, len(cache_content))
self.assertTrue(cache_content[1].endswith(".descriptor.txt"))
module_files = sorted(tf.gfile.ListDirectory(
os.path.join(cache_dir, cache_content[0])))
self.assertListEqual(["saved_model.pb", "tfhub_module.pb"], module_files)
finally:
os.unsetenv("TFHUB_CACHE_DIR")
def _model_fn(features, labels, mode):
"""A model_fn that uses a mock TF-Hub module."""
del labels
spec = hub.create_module_spec(text_module_fn)
embedding = hub.Module(spec)
if register_module:
hub.register_module_for_export(embedding, _EXPORT_MODULE_NAME)
predictions = embedding(features[_TEXT_FEATURE_NAME])
loss = tf.constant(0.0)
global_step = tf.train.get_global_step()
train_op = tf.assign_add(global_step, 1)
return tf.estimator.EstimatorSpec(
mode=mode,
predictions=predictions,
loss=loss,
train_op=train_op)
def testModuleWithTrainedVariable(self):
export_path = os.path.join(self.get_temp_dir(), "var-module")
with tf.Graph().as_default():
spec = hub.create_module_spec(stateful_module_fn)
m = hub.Module(spec, trainable=True)
assign_op = tf.assign(m.variable_map["var123"],
tf.constant([9.0, 9.0, 9.0]))
with tf.Session() as sess:
sess.run(assign_op)
m.export(export_path, sess)
with tf.Graph().as_default():
f = hub.Module(export_path)
out = f()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
got = sess.run(out)
self.assertAllClose(got, [9.0, 9.0, 9.0])
def testSparseTensors(self):
square_spec = hub.create_module_spec(sparse_square_module_fn)
with tf.Graph().as_default():
square = hub.Module(square_spec)
v = tf.sparse_placeholder(dtype=tf.int64, name="v")
y = square(v)
with tf.Session().as_default():
indices = [[0, 0], [0, 1], [1, 1]]
values = [10, 2, 1]
shape = [2, 2]
v1 = tf.SparseTensorValue(indices, values, shape)
v2 = y.eval(feed_dict={v: v1})
v4 = y.eval(feed_dict={v: v2})
self.assertAllEqual(v4.indices, indices) # Unchanged.
self.assertAllEqual(v4.values, [t**4 for t in values]) # Squared twice.
self.assertAllEqual(v4.dense_shape, shape) # Unchanged.
