Function _create_frozen_model

tfjs-converter/python/test_pip_package.py:148–186

(save_path)

    m.export(save_path, sess)

def _create_frozen_model(save_path):
  graph = tf.Graph()
  saved_model_dir = os.path.join(save_path)
  with graph.as_default():
    x = tf.constant([[37.0, -23.0], [1.0, 4.0]])
    w = tf.Variable(tf.random.uniform([2, 2]))
    y = tf.matmul(x, w)
    tf.nn.softmax(y)
    init_op = w.initializer

    # Create a builder
    builder = saved_model.builder.SavedModelBuilder(
        saved_model_dir)

    with tf.compat.v1.Session() as sess:
      # Run the initializer on `w`.
      sess.run(init_op)

      builder.add_meta_graph_and_variables(
          sess, [saved_model.tag_constants.SERVING],
          signature_def_map=None,
          assets_collection=None)

    builder.save()

  frozen_file = os.path.join(save_path, 'frozen.pb')
  freeze_graph.freeze_graph(
      '',
      '',
      True,
      '',
      "Softmax",
      '',
      '',
      frozen_file,
      True,
      '',
      saved_model_tags=saved_model.tag_constants.SERVING,
      input_saved_model_dir=saved_model_dir)
class APIAndShellTest(tf.test.TestCase):
  """Tests for the Python API of the pip package."""