Function _createTensorFlowSavedModelV1

tfjs-converter/python/test_pip_package.py:76–109

Create a TensorFlow SavedModel for testing. Args: name_scope: Name scope to create the model under. This helps avoid op and variable name clashes between different test methods. save_path: The directory path in which to save the model.

(name_scope, save_path)

  return model

def _createTensorFlowSavedModelV1(name_scope, save_path):
  """Create a TensorFlow SavedModel for testing.
  Args:
    name_scope: Name scope to create the model under. This helps avoid
      op and variable name clashes between different test methods.
    save_path: The directory path in which to save the model.
  """
  graph = tf.Graph()
  with graph.as_default():
    with tf.compat.v1.name_scope(name_scope):
      x = tf.compat.v1.constant([[37.0, -23.0], [1.0, 4.0]])
      w = tf.compat.v1.get_variable('w', shape=[2, 2])
      y = tf.compat.v1.matmul(x, w)
      output = tf.compat.v1.nn.softmax(y)
      init_op = w.initializer

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

      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={
                "serving_default":
                    saved_model.signature_def_utils.predict_signature_def(
                        inputs={"x": x},
                        outputs={"output": output})
            },
            assets_collection=None)

      builder.save()

def _createTensorFlowSavedModel(save_path):
  """Create a TensorFlow SavedModel for testing.