Python tensorflow_hub.add_signature() Examples

def layers_module_fn():
  """Module that exercises the use of layers."""
  # This is a plain linear map Mx+b regularized by the sum of the squares
  # of the coefficients in M and b.
  x = tf_v1.placeholder(dtype=tf.float32, shape=[None, 2], name="x")
  def l2(weights):
    """Applies l2 regularization to weights."""
    with tf.control_dependencies([weights]):
      return 2.0 * tf_v1.nn.l2_loss(weights)

  h = tf_v1.layers.dense(
      x, 2,
      activation=None,
      kernel_regularizer=l2,
      bias_regularizer=l2)
  hub.add_signature(inputs=x, outputs=h) 

def apply_model(image_fn,  # pylint: disable=missing-docstring
                is_training,
                num_outputs,
                make_signature=False):

  # Image tensor needs to be created lazily in order to satisfy tf-hub
  # restriction: all tensors should be created inside tf-hub helper function.
  images = image_fn()

  net = get_net(num_classes=num_outputs)

  output, end_points = net(images, is_training)

  if make_signature:
    hub.add_signature(inputs={'image': images}, outputs=output)
    hub.add_signature(
        name='representation',
        inputs={'image': images},
        outputs=end_points)
  return output 

def apply_model(image_fn,  # pylint: disable=missing-docstring
                is_training,
                num_outputs,
                make_signature=False):

  # Image tensor needs to be created lazily in order to satisfy tf-hub
  # restriction: all tensors should be created inside tf-hub helper function.
  images = image_fn()

  net = get_net(num_classes=num_outputs)

  output, end_points = net(images, is_training)

  if make_signature:
    hub.add_signature(inputs={'image': images}, outputs=output)
    hub.add_signature(inputs={'image': images}, outputs=end_points,
                      name='representation')
  return output 

def _save_half_plus_one_hub_module_v1(path):
  """Writes TF1.x hub.Module to compute y = wx + 1, with w trainable."""
  def half_plus_one():
    x = tf.compat.v1.placeholder(shape=(None,1), dtype=tf.float32)
    # Use TF1 native tf.compat.v1.layers instead of tf.keras.layers as they
    # correctly update TF collections, such as REGULARIZATION_LOSS.
    times_w = tf.compat.v1.layers.Dense(
        units=1,
        kernel_initializer=tf.keras.initializers.Constant([[0.5]]),
        kernel_regularizer=tf.keras.regularizers.l2(0.01),
        use_bias=False)
    plus_1 = tf.compat.v1.layers.Dense(
        units=1,
        kernel_initializer=tf.keras.initializers.Constant([[1.0]]),
        bias_initializer=tf.keras.initializers.Constant([1.0]),
        trainable=False)
    y = plus_1(times_w(x))
    hub.add_signature(inputs=x, outputs=y)

  spec = hub.create_module_spec(half_plus_one)
  _export_module_spec_with_init_weights(spec, path) 

def text_module_fn():
  embeddings = [
      ("", [0, 0, 0, 0]),  # OOV items are mapped to this embedding.
      ("hello world", [1, 2, 3, 4]),
      ("pair-programming", [5, 5, 5, 5]),
  ]
  keys = tf.constant([item[0] for item in embeddings], dtype=tf.string)
  indices = tf.constant(list(range(len(embeddings))), dtype=tf.int64)
  tbl_init = KeyValueTensorInitializer(keys, indices)
  table = HashTable(tbl_init, 0)

  weights_initializer = tf.cast(
      tf.constant(list([item[1] for item in embeddings])), tf.float32)

  weights = tf_v1.get_variable(
      "weights", dtype=tf.float32, initializer=weights_initializer)

  text_tensor = tf_v1.placeholder(dtype=tf.string, name="text", shape=[None])
  indices_tensor = table.lookup(text_tensor)
  embedding_tensor = tf.gather(weights, indices_tensor)
  hub.add_signature(inputs=text_tensor, outputs=embedding_tensor) 

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_v1.placeholder(dtype=tf.int64, name="indices1")
      indices2 = tf_v1.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_v1.Session() as sess:
        sess.run(tf_v1.tables_initializer())
        module_a.export(export_path, sess) 

def valid_colocation_module_fn():
  w = tf.Variable(42 + 69, name="w")
  # w.op has the same name on resource and non-resource variables
  with tf_v1.colocate_with(w.op):
    # A colocation reference among state nodes is ok.
    v = tf.Variable(1.0, name="v")
    assert v.op.colocation_groups() == [tf.compat.as_bytes("loc:@w")]
    # A colocation reference from other nodes to state nodes is ok.
    y = tf.add(v, 1, name="y")
    assert y.op.colocation_groups() == [tf.compat.as_bytes("loc:@w")]
  x = tf_v1.placeholder(dtype=tf.float32, name="x")
  with tf_v1.colocate_with(x):
    # A colocation reference from other nodes to input nodes is ok.
    z = tf.add(x, 1, name="z")
    assert z.op.colocation_groups() == [tf.compat.as_bytes("loc:@x")]
  hub.add_signature(inputs=dict(x=x), outputs=dict(y=y, z=z)) 

def stateful_rv_with_input_module_fn():
  x = tf_v1.placeholder(dtype=tf.float32, name="x")
  # Add a placeholder/variable that doesn't go to an output.
  y = tf_v1.placeholder(dtype=tf.float32, name="y")
  r = tf_v1.get_variable(
      "rv_var123",
      shape=[],
      initializer=tf_v1.constant_initializer(10.0),
      use_resource=True)
  t = tf_v1.get_variable(
      "rv_var456",
      shape=[],
      initializer=tf_v1.constant_initializer(10.0),
      use_resource=True)
  t.assign(y)
  res = x + r
  hub.add_signature(inputs={"x": x}, outputs=res) 

def testUnsupportedCollections(self):

    def module_fn():
      scale = tf_v1.get_variable("x", (), collections=["my_scope"])
      x = tf_v1.placeholder(tf.float32, shape=[None, 3])
      native_module.add_signature("my_func", {"x": x}, {"y": x*scale})

    with self.assertRaises(ValueError) as cm:
      _ = native_module.create_module_spec(module_fn)
      self.assertIn("Unsupported collections in graph", cm)

    with tf.Graph().as_default() as tmp_graph:
      module_fn()
      unsupported_collections = native_module.get_unsupported_collections(
          tmp_graph.get_all_collection_keys())
      self.assertEqual(["my_scope"], unsupported_collections)

    _ = native_module.create_module_spec(
        module_fn, drop_collections=unsupported_collections) 

def nested_cond_module_fn():
  """Computes relu(x) with nested conditionals."""
  x = tf_v1.placeholder(dtype=tf.float32, name="x", shape=[])
  # pylint: disable=g-long-lambda
  result = tf.cond(
      0 < x,
      lambda: tf.cond(3 < x,
                      lambda: tf.identity(x),
                      lambda: tf.multiply(x, 1.0)),
      lambda: tf.cond(x < -3,
                      lambda: tf.constant(0.0),
                      lambda: tf.multiply(0.0, 1.0)))
  # pylint: enable=g-long-lambda
  hub.add_signature(inputs=x, outputs=result) 

def brittle_multivalued_colocation_module_fn():
  x, y = tf.split([1, 2], 2, name="split")
  with tf_v1.colocate_with(x), tf_v1.colocate_with(y):
    z = tf.add(x, y, name="add")
    assert z.op.colocation_groups() == [tf.compat.as_bytes("loc:@split")]
  hub.add_signature(inputs=dict(x=x, y=y), outputs=z, name="both")
  hub.add_signature(inputs=dict(x=x), outputs=z, name="partial") 

def update_ops_module_fn():
  counter = tf.Variable(0, trainable=False)
  tf_v1.add_to_collection(tf_v1.GraphKeys.UPDATE_OPS, counter.assign_add(1))
  hub.add_signature(inputs=None, outputs=counter.value()) 

def batch_norm_module_fn(is_training):
  """Module that exercises batch normalization, incl. UPDATE_OPS."""
  x = tf_v1.placeholder(dtype=tf.float32, shape=[None, 1], name="x")
  y = tf_v1.layers.batch_normalization(
      momentum=0.4,
      inputs=x,
      fused=False,
      training=is_training)
  hub.add_signature(inputs=x, outputs=y) 

def multiple_outputs_module_fn():
  x = tf_v1.placeholder(dtype=tf.float32)
  v = tf.Variable([3.0])
  hub.add_signature(
      inputs={"x": x},
      outputs={"y": v * x, "z": v * v * x}) 

def create_assets_module_fn(vocabulary_file):

  def assets_module_fn():
    indices = tf_v1.placeholder(dtype=tf.int64, name="indices")
    outputs = do_table_lookup(indices, vocabulary_file)
    hub.add_signature(inputs=indices, outputs=outputs)

  return assets_module_fn 

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 another_stateful_module_fn():
  """Stateful module with inputs."""
  module_input = tf_v1.placeholder(dtype=tf.float32)
  variable = tf.Variable([3.0], name="iamtheoneandonly")
  hub.add_signature(inputs=module_input, outputs=module_input*variable) 

def multiple_signature_module_fn():
  """Stateful module with multiple signatures."""
  weight = tf.Variable([3.0])

  x_input = tf_v1.placeholder(dtype=tf.float32)
  x_output = tf.multiply(x_input, weight)
  hub.add_signature("mul", inputs=x_input, outputs=x_output)

  y_input = tf_v1.placeholder(dtype=tf.float32)
  y_output = tf.divide(y_input, weight)
  hub.add_signature("div", inputs=y_input, outputs=y_output) 

def cond_module_fn():
  """Computes relu(x) with a conditional."""
  x = tf_v1.placeholder(dtype=tf.float32, name="x", shape=[])
  result = tf.cond(0 < x, lambda: tf.identity(x), lambda: tf.constant(0.0))
  hub.add_signature(inputs=x, outputs=result) 

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 nested_control_flow_module_fn():
  """Compute the sum of elements greater than 'a' with nested control flow."""
  elems = tf_v1.placeholder(
      dtype=tf.float32, name="elems", shape=[None])
  a = tf_v1.placeholder(dtype=tf.float32, name="a")

  def sum_above_a(acc, x):
    return acc + tf.cond(x > a, lambda: x, lambda: 0.0)

  hub.add_signature(
      inputs={"elems": elems, "a": a},
      outputs=tf.foldl(sum_above_a, elems, initializer=tf.constant(0.0))) 

def attached_messages_module_fn(tagged=0):
  x = tf_v1.placeholder(tf.float32, shape=[None])
  hub.add_signature(inputs={"x": x}, outputs={"y": 2*x})
  # For brevity, this test borrows two well-known, stable message types
  # from TensorFlow. They are not likely choices for actual uses.
  hub.attach_message("numbers",
                     tf_v1.train.Int64List(value=[-3]))  # Overwritten.
  hub.attach_message("numbers", tf_v1.train.Int64List(value=[42, 69]))
  hub.attach_message("letters", tf_v1.train.BytesList(value=[
      tf.compat.as_bytes("abc"), tf.compat.as_bytes("xyz")]))
  hub.attach_message("tagged", tf_v1.train.Int64List(value=[tagged])) 

def module_fn(self, dim=10):
    x = tf_v1.placeholder(dtype=tf.float32, shape=[None, dim])
    y = self.f(x, dim=dim)
    hub.add_signature(inputs=x, outputs=y) 

def invalid_text_module_fn():
  text = tf_v1.placeholder(tf.string, shape=[10])
  hub.add_signature(inputs=text, outputs=tf.zeros([10, 3])) 

def text_module_fn():
  weights = tf_v1.get_variable(
      "weights", dtype=tf.float32, shape=[100, 10])
  #      initializer=tf.random_uniform_initializer())
  text = tf_v1.placeholder(tf.string, shape=[None])
  hash_buckets = tf_v1.string_to_hash_bucket_fast(text, weights.get_shape()[0])
  embeddings = tf_v1.gather(weights, hash_buckets)
  hub.add_signature(inputs=text, outputs=embeddings) 

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 half_plus_two():
  a = tf.get_variable("a", shape=[])
  b = tf.get_variable("b", shape=[])
  x = tf.placeholder(tf.float32)
  y = a*x + b
  hub.add_signature(inputs=x, outputs=y) 

def apply_model(image_fn,
                is_training,
                num_outputs,
                make_signature=False):
  """Creates the patch based model output from patches representations.

  Args:
    image_fn: function returns image tensor.
    is_training: is training flag used for batch norm and drop out.
    num_outputs: number of output classes.
    make_signature: whether to create signature for hub module.


  Returns:
    out: output tensor with shape [n*m, 1, 1, num_outputs].

  Raises:
    ValueError: An error occurred when the architecture is unknown.
  """
  # Image tensor needs to be created lazily in order to satisfy tf-hub
  # restriction: all tensors should be created inside tf-hub helper function.
  images = image_fn()

  net = get_net(num_classes=num_outputs)
  out, end_points = net(images, is_training,
                        weight_decay=FLAGS.get_flag_value('weight_decay', 1e-4))

  print(end_points)

  if len(out.get_shape().as_list()) == 4:
    out = tf.squeeze(out, [1, 2])

  if make_signature:
    hub.add_signature(inputs={'image': images}, outputs=out)
    hub.add_signature(
        name='representation',
        inputs={'image': images},
        outputs=end_points)
  return out 

def stateful_module_fn():
  v = tf_v1.get_variable(
      "var123", shape=[3],
      initializer=tf_v1.constant_initializer([1.0, 2.0, 3.0]))
  hub.add_signature(outputs=v.value()) 

def multi_signature_module():
  x = tf_v1.placeholder(tf.float32, shape=[None])
  native_module.add_signature("double", {"x": x}, {"y": 2*x})

  z = tf_v1.placeholder(tf.float32, shape=[None])
  native_module.add_signature("square", {"z": z}, {"z_out": z*z})