Support different label type of tfrecords

Author: tobegit3hub

dense_classifier.py

@@ -16,6 +16,11 @@
 import util
 import model
 
+logging.basicConfig(
+    format='%(asctime)s %(levelname)-8s %(message)s',
+    level=logging.INFO,
+    datefmt='%Y-%m-%d %H:%M:%S')
+
 
 def define_flags():
   """
@@ -89,14 +94,17 @@ def define_flags():
   ])
 
   # Print flags
+  FLAGS.mode
   parameter_value_map = {}
   for key in FLAGS.__flags.keys():
     parameter_value_map[key] = FLAGS.__flags[key].value
   pprint.PrettyPrinter().pprint(parameter_value_map)
-
   return FLAGS
 
 
+FLAGS = define_flags()
+
+
 def parse_tfrecords_function(example_proto):
   """
   Decode TFRecords for Dataset.
@@ -175,10 +183,6 @@ def inference(inputs, input_units, output_units, is_train=True):
                                FLAGS)
 
 
-logging.basicConfig(level=logging.INFO)
-FLAGS = define_flags()
-
-
 def main():
   """
   Train the TensorFlow models.

sparse_classifier.py

@@ -14,9 +14,13 @@
     signature_constants, signature_def_utils, tag_constants, utils)
 
 import sparse_model
-import model
 import util
 
+logging.basicConfig(
+    format='%(asctime)s %(levelname)-8s %(message)s',
+    level=logging.INFO,
+    datefmt='%Y-%m-%d %H:%M:%S')
+
 
 def define_flags():
   """
@@ -34,6 +38,7 @@ def define_flags():
                       "The glob pattern of train TFRecords files")
   flags.DEFINE_integer("feature_size", 124, "Number of feature size")
   flags.DEFINE_integer("label_size", 2, "Number of label size")
+  flags.DEFINE_string("label_type", "int", "The type of label")
   flags.DEFINE_float("learning_rate", 0.01, "The learning rate")
   flags.DEFINE_integer("epoch_number", 10, "Number of epochs to train")
   flags.DEFINE_integer("batch_size", 1024, "The batch size of training")
@@ -81,14 +86,17 @@ def define_flags():
   ])
 
   # Print flags
+  FLAGS.mode
   parameter_value_map = {}
   for key in FLAGS.__flags.keys():
     parameter_value_map[key] = FLAGS.__flags[key].value
   pprint.PrettyPrinter().pprint(parameter_value_map)
-
   return FLAGS
 
 
+FLAGS = define_flags()
+
+
 def parse_tfrecords_function(example_proto):
   """
     Decode TFRecords for Dataset.
@@ -100,15 +108,31 @@ def parse_tfrecords_function(example_proto):
       The op of features and labels
     """
 
-  features = {
-      "ids": tf.VarLenFeature(tf.int64),
-      "values": tf.VarLenFeature(tf.float32),
-      "label": tf.FixedLenFeature([], tf.int64, default_value=0)
-  }
+  if FLAGS.label_type == "int":
+    features = {
+        "ids": tf.VarLenFeature(tf.int64),
+        "values": tf.VarLenFeature(tf.float32),
+        "label": tf.FixedLenFeature([], tf.int64, default_value=0)
+    }
 
-  parsed_features = tf.parse_single_example(example_proto, features)
-  return parsed_features["label"], parsed_features["ids"], parsed_features[
-      "values"]
+    parsed_features = tf.parse_single_example(example_proto, features)
+    labels = parsed_features["label"]
+    ids = parsed_features["ids"]
+    values = parsed_features["values"]
+
+  elif FLAGS.label_type == "float":
+    features = {
+        "ids": tf.VarLenFeature(tf.int64),
+        "values": tf.VarLenFeature(tf.float32),
+        "label": tf.FixedLenFeature([], tf.float32, default_value=0)
+    }
+
+    parsed_features = tf.parse_single_example(example_proto, features)
+    labels = tf.cast(parsed_features["label"], tf.int32)
+    ids = parsed_features["ids"]
+    values = parsed_features["values"]
+
+  return labels, ids, values
 
 
 def inference(sparse_ids, sparse_values, is_train=True):
@@ -133,10 +157,6 @@ def inference(sparse_ids, sparse_values, is_train=True):
                                              is_train, FLAGS)
 
 
-logging.basicConfig(level=logging.INFO)
-FLAGS = define_flags()
-
-
 def main():
 
   if os.path.exists(FLAGS.checkpoint_path) == False:
@@ -170,8 +190,8 @@ def main():
   validation_filename_placeholder = tf.placeholder(tf.string, shape=[None])
   validation_dataset = tf.data.TFRecordDataset(validation_filename_placeholder)
   validation_dataset = validation_dataset.map(parse_tfrecords_function).repeat(
-      epoch_number).batch(FLAGS.validation_batch_size).shuffle(
-          buffer_size=validation_buffer_size)
+  ).batch(FLAGS.validation_batch_size).shuffle(
+      buffer_size=validation_buffer_size)
   validation_dataset_iterator = validation_dataset.make_initializable_iterator(
   )
   validation_labels, validation_ids, validation_values = validation_dataset_iterator.get_next(

util.py

@@ -1,16 +1,14 @@
-
 from __future__ import absolute_import, division, print_function
 
 import logging
 import os
 import tensorflow as tf
 from tensorflow.python.saved_model import builder as saved_model_builder
-from tensorflow.python.saved_model import (
-    signature_constants, tag_constants)
+from tensorflow.python.saved_model import (signature_constants, tag_constants)
 
 
 def get_optimizer_by_name(optimizer_name, learning_rate):
-    """
+  """
     Get optimizer object by the optimizer name.
     
     Args:
@@ -21,30 +19,30 @@ def get_optimizer_by_name(optimizer_name, learning_rate):
       The optimizer object.
     """
 
-    logging.info("Use the optimizer: {}".format(optimizer_name))
-    if optimizer_name == "sgd":
-        optimizer = tf.train.GradientDescentOptimizer(learning_rate)
-    elif optimizer_name == "adadelta":
-        optimizer = tf.train.AdadeltaOptimizer(learning_rate)
-    elif optimizer_name == "adagrad":
-        optimizer = tf.train.AdagradOptimizer(learning_rate)
-    elif optimizer_name == "adam":
-        optimizer = tf.train.AdamOptimizer(learning_rate)
-    elif optimizer_name == "ftrl":
-        optimizer = tf.train.FtrlOptimizer(learning_rate)
-    elif optimizer_name == "rmsprop":
-        optimizer = tf.train.RMSPropOptimizer(learning_rate)
-    else:
-        optimizer = tf.train.GradientDescentOptimizer(learning_rate)
-    return optimizer
+  logging.info("Use the optimizer: {}".format(optimizer_name))
+  if optimizer_name == "sgd":
+    optimizer = tf.train.GradientDescentOptimizer(learning_rate)
+  elif optimizer_name == "adadelta":
+    optimizer = tf.train.AdadeltaOptimizer(learning_rate)
+  elif optimizer_name == "adagrad":
+    optimizer = tf.train.AdagradOptimizer(learning_rate)
+  elif optimizer_name == "adam":
+    optimizer = tf.train.AdamOptimizer(learning_rate)
+  elif optimizer_name == "ftrl":
+    optimizer = tf.train.FtrlOptimizer(learning_rate)
+  elif optimizer_name == "rmsprop":
+    optimizer = tf.train.RMSPropOptimizer(learning_rate)
+  else:
+    optimizer = tf.train.GradientDescentOptimizer(learning_rate)
+  return optimizer
 
 
 def save_model(model_path,
-        model_version,
-        sess,
-        signature_def_map,
-        is_save_graph=False):
-    """
+               model_version,
+               sess,
+               signature_def_map,
+               is_save_graph=False):
+  """
     Save the model in standard SavedModel format.
     
     Args:
@@ -58,36 +56,36 @@ def save_model(model_path,
       None
     """
 
-    export_path = os.path.join(model_path, str(model_version))
-    if os.path.isdir(export_path) == True:
-        logging.error("The model exists in path: {}".format(export_path))
-        return
+  export_path = os.path.join(model_path, str(model_version))
+  if os.path.isdir(export_path) == True:
+    logging.error("The model exists in path: {}".format(export_path))
+    return
 
-    try:
-        # Save the SavedModel
-        legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
-        builder = saved_model_builder.SavedModelBuilder(export_path)
-        builder.add_meta_graph_and_variables(
-                sess, [tag_constants.SERVING],
-                clear_devices=True,
-                signature_def_map=signature_def_map,
-                legacy_init_op=legacy_init_op)
-        logging.info("Save the model in: {}".format(export_path))
-        builder.save()
+  try:
+    # Save the SavedModel
+    legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
+    builder = saved_model_builder.SavedModelBuilder(export_path)
+    builder.add_meta_graph_and_variables(
+        sess, [tag_constants.SERVING],
+        clear_devices=True,
+        signature_def_map=signature_def_map,
+        legacy_init_op=legacy_init_op)
+    logging.info("Save the model in: {}".format(export_path))
+    builder.save()
 
-        # Save the GraphDef
-        if is_save_graph == True:
-            graph_file_name = "graph.pb"
-            logging.info("Save the graph file in: {}".format(model_path))
-            tf.train.write_graph(
-                    sess.graph_def, model_path, graph_file_name, as_text=False)
+    # Save the GraphDef
+    if is_save_graph == True:
+      graph_file_name = "graph.pb"
+      logging.info("Save the graph file in: {}".format(model_path))
+      tf.train.write_graph(
+          sess.graph_def, model_path, graph_file_name, as_text=False)
 
-    except Exception as e:
-        logging.error("Fail to export saved model, exception: {}".format(e))
+  except Exception as e:
+    logging.error("Fail to export saved model, exception: {}".format(e))
 
 
 def restore_from_checkpoint(sess, saver, checkpoint_file_path):
-    """
+  """
     Restore session from checkpoint files.
     
     Args:
@@ -98,11 +96,11 @@ def restore_from_checkpoint(sess, saver, checkpoint_file_path):
     Return:
       True if restore successfully and False if fail
     """
-    if checkpoint_file_path:
-        logging.info(
-                "Restore session from checkpoint: {}".format(checkpoint_file_path))
-        saver.restore(sess, checkpoint_file_path)
-        return True
-    else:
-        logging.error("Checkpoint not found: {}".format(checkpoint_file_path))
-        return False
+  if checkpoint_file_path:
+    logging.info(
+        "Restore session from checkpoint: {}".format(checkpoint_file_path))
+    saver.restore(sess, checkpoint_file_path)
+    return True
+  else:
+    logging.error("Checkpoint not found: {}".format(checkpoint_file_path))
+    return False