Python tensorflow_hub.create_module_spec函数代码示例

  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 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 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 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 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 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 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 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 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 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 _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 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])