Define/export a single-node TF graph for inferencing

def export_fun(args):
  """Define/export a single-node TF graph for inferencing"""
  # Input placeholder for inferencing
  x = tf.placeholder(tf.float32, [None, IMAGE_PIXELS * IMAGE_PIXELS], name="x")

  # Variables of the hidden layer
  hid_w = tf.Variable(tf.truncated_normal([IMAGE_PIXELS * IMAGE_PIXELS, hidden_units],
                      stddev=1.0 / IMAGE_PIXELS), name="hid_w")
  hid_b = tf.Variable(tf.zeros([hidden_units]), name="hid_b")
  tf.summary.histogram("hidden_weights", hid_w)

  # Variables of the softmax layer
  sm_w = tf.Variable(tf.truncated_normal([hidden_units, 10],
                     stddev=1.0 / math.sqrt(hidden_units)), name="sm_w")
  sm_b = tf.Variable(tf.zeros([10]), name="sm_b")

  hid_lin = tf.nn.xw_plus_b(x, hid_w, hid_b)
  hid = tf.nn.relu(hid_lin)
  y = tf.nn.softmax(tf.nn.xw_plus_b(hid, sm_w, sm_b))
  prediction = tf.argmax(y, 1, name="prediction")

  saver = tf.train.Saver()

  with tf.Session() as sess:
    # load graph from a checkpoint
    logging.info("model path: {}".format(args.model_dir))
    ckpt = tf.train.get_checkpoint_state(args.model_dir)
    logging.info("ckpt: {}".format(ckpt))
    assert ckpt and ckpt.model_checkpoint_path, "Invalid model checkpoint path: {}".format(args.model_dir)
    saver.restore(sess, ckpt.model_checkpoint_path)

    logging.info("Exporting saved_model to: {}".format(args.export_dir))
    # exported signatures defined in code
    signatures = {
      tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: {
        'inputs': {'image': x},
        'outputs': {'prediction': prediction},
        'method_name': tf.saved_model.signature_constants.PREDICT_METHOD_NAME
      },
      'featurize': {
        'inputs': {'image': x},
        'outputs': {'features': hid},
        'method_name': 'featurize'
      }
    }
    TFNode.export_saved_model(sess,
                              args.export_dir,
                              tf.saved_model.tag_constants.SERVING,
                              signatures)
    logging.info("Exported saved_model")