def testCustomMainOp(self): export_dir = os.path.join(test.get_temp_dir(), "test_main_op") builder = saved_model_builder.SavedModelBuilder(export_dir) with self.test_session(graph=ops.Graph()) as sess: # Add `v1` and `v2` variables to the graph. v1 = variables.Variable(1, name="v1") ops.add_to_collection("v", v1) v2 = variables.Variable(2, name="v2") ops.add_to_collection("v", v2) # Initialize another variable `v3` to 42. v3 = variables.Variable(42, name="v3") ops.add_to_collection("v", v3) # Set up an assignment op to be run as part of the main_op. with ops.control_dependencies([main_op.main_op()]): add_v1_v2 = math_ops.add(v1._ref(), v2._ref()) custom_main_op = control_flow_ops.group( state_ops.assign(v3, add_v1_v2)) sess.run(custom_main_op) builder.add_meta_graph_and_variables(sess, ["foo"], main_op=custom_main_op) # Save the SavedModel to disk. builder.save() with self.test_session(graph=ops.Graph()) as sess: loader.load(sess, ["foo"], export_dir) self.assertEqual(1, ops.get_collection("v")[0].eval()) self.assertEqual(2, ops.get_collection("v")[1].eval()) # Evaluates to the sum of the first two variables and assigned as part of # the main_op, following a restore. self.assertEqual(3, ops.get_collection("v")[2].eval())
def testCustomMainOp(self): export_dir = self._get_export_dir("test_main_op") builder = saved_model_builder.SavedModelBuilder(export_dir) with self.test_session(graph=ops.Graph()) as sess: # Add `v1` and `v2` variables to the graph. v1 = variables.Variable(1, name="v1") ops.add_to_collection("v", v1) v2 = variables.Variable(2, name="v2") ops.add_to_collection("v", v2) # Initialize another variable `v3` to 42. v3 = variables.Variable(42, name="v3") ops.add_to_collection("v", v3) # Set up an assignment op to be run as part of the main_op. with ops.control_dependencies([main_op.main_op()]): add_v1_v2 = math_ops.add(v1._ref(), v2._ref()) custom_main_op = control_flow_ops.group(state_ops.assign(v3, add_v1_v2)) sess.run(custom_main_op) builder.add_meta_graph_and_variables( sess, ["foo"], main_op=custom_main_op) # Save the SavedModel to disk. builder.save() with self.test_session(graph=ops.Graph()) as sess: loader.load(sess, ["foo"], export_dir) self.assertEqual(1, ops.get_collection("v")[0].eval()) self.assertEqual(2, ops.get_collection("v")[1].eval()) # Evaluates to the sum of the first two variables and assigned as part of # the main_op, following a restore. self.assertEqual(3, ops.get_collection("v")[2].eval())
def testCustomMainOp(self): export_dir = os.path.join(tf.test.get_temp_dir(), "test_main_op") builder = saved_model_builder.SavedModelBuilder(export_dir) with self.test_session(graph=tf.Graph()) as sess: # Add `v1` and `v2` variables to the graph. v1 = tf.Variable(1, name="v1") tf.add_to_collection("v", v1) v2 = tf.Variable(2, name="v2") tf.add_to_collection("v", v2) # Initialize another variable `v3` to 42. v3 = tf.Variable(42, name="v3", trainable=False, collections=[]) tf.add_to_collection("v", v3) # Set up an assignment op to be run as part of the main_op. assign_v3 = tf.assign(v3, tf.add(v1, v2)) custom_main_op = tf.group(main_op.main_op(), assign_v3) sess.run(tf.global_variables_initializer()) builder.add_meta_graph_and_variables(sess, ["foo"], main_op=custom_main_op) # Save the SavedModel to disk. builder.save() with self.test_session(graph=tf.Graph()) as sess: loader.load(sess, ["foo"], export_dir) self.assertEqual(1, tf.get_collection("v")[0].eval()) self.assertEqual(2, tf.get_collection("v")[1].eval()) # Evaluates to the sum of the first two variables and assigned as part of # the main_op, following a restore. self.assertEqual(3, tf.get_collection("v")[2].eval())
def _generate_saved_model_for_half_plus_two(export_dir, as_text=False, use_main_op=False): """Generates SavedModel for half plus two. Args: export_dir: The directory to which the SavedModel should be written. as_text: Writes the SavedModel protocol buffer in text format to disk. use_main_op: Whether to supply a main op during SavedModel build time. """ builder = tf.saved_model.builder.SavedModelBuilder(export_dir) with tf.Session(graph=tf.Graph()) as sess: # Set up the model parameters as variables to exercise variable loading # functionality upon restore. a = tf.Variable(0.5, name="a") b = tf.Variable(2.0, name="b") c = tf.Variable(3.0, name="c") # Create a placeholder for serialized tensorflow.Example messages to be fed. serialized_tf_example = tf.placeholder(tf.string, name="tf_example") # Parse the tensorflow.Example looking for a feature named "x" with a single # floating point value. feature_configs = { "x": tf.FixedLenFeature( [1], dtype=tf.float32), "x2": tf.FixedLenFeature( [1], dtype=tf.float32, default_value=[0.0]) } tf_example = tf.parse_example(serialized_tf_example, feature_configs) # Use tf.identity() to assign name x = tf.identity(tf_example["x"], name="x") y = tf.add(tf.multiply(a, x), b, name="y") y2 = tf.add(tf.multiply(a, x), c, name="y2") x2 = tf.identity(tf_example["x2"], name="x2") y3 = tf.add(tf.multiply(a, x2), c, name="y3") # Create an assets file that can be saved and restored as part of the # SavedModel. original_assets_directory = "/tmp/original/export/assets" original_assets_filename = "foo.txt" original_assets_filepath = _write_assets(original_assets_directory, original_assets_filename) # Set up the assets collection. assets_filepath = tf.constant(original_assets_filepath) tf.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, assets_filepath) filename_tensor = tf.Variable( original_assets_filename, name="filename_tensor", trainable=False, collections=[]) assign_filename_op = filename_tensor.assign(original_assets_filename) # Set up the signature for Predict with input and output tensor # specification. predict_input_tensor = tf.saved_model.utils.build_tensor_info(x) predict_signature_inputs = {"x": predict_input_tensor} predict_output_tensor = tf.saved_model.utils.build_tensor_info(y) predict_signature_outputs = {"y": predict_output_tensor} predict_signature_def = ( tf.saved_model.signature_def_utils.build_signature_def( predict_signature_inputs, predict_signature_outputs, tf.saved_model.signature_constants.PREDICT_METHOD_NAME)) signature_def_map = { "regress_x_to_y": _build_regression_signature(serialized_tf_example, y), "regress_x_to_y2": _build_regression_signature(serialized_tf_example, y2), "regress_x2_to_y3": _build_regression_signature(x2, y3), "classify_x_to_y": _build_classification_signature(serialized_tf_example, y), tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: predict_signature_def } # Initialize all variables and then save the SavedModel. sess.run(tf.global_variables_initializer()) signature_def_map = { "regress_x_to_y": _build_regression_signature(serialized_tf_example, y), "regress_x_to_y2": _build_regression_signature(serialized_tf_example, y2), "regress_x2_to_y3": _build_regression_signature(x2, y3), "classify_x_to_y": _build_classification_signature(serialized_tf_example, y), "classify_x2_to_y3": _build_classification_signature(x2, y3), tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: predict_signature_def } if use_main_op: builder.add_meta_graph_and_variables( sess, [tf.saved_model.tag_constants.SERVING], signature_def_map=signature_def_map, assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS), main_op=tf.group(main_op.main_op(), assign_filename_op)) else: builder.add_meta_graph_and_variables( sess, [tf.saved_model.tag_constants.SERVING], signature_def_map=signature_def_map, assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS), legacy_init_op=tf.group(assign_filename_op)) builder.save(as_text)
def _generate_saved_model_for_half_plus_two(export_dir, as_text=False, use_main_op=False): """Generates SavedModel for half plus two. Args: export_dir: The directory to which the SavedModel should be written. as_text: Writes the SavedModel protocol buffer in text format to disk. use_main_op: Whether to supply a main op during SavedModel build time. """ builder = tf.saved_model.builder.SavedModelBuilder(export_dir) with tf.Session(graph=tf.Graph()) as sess: # Set up the model parameters as variables to exercise variable loading # functionality upon restore. a = tf.Variable(0.5, name="a") b = tf.Variable(2.0, name="b") c = tf.Variable(3.0, name="c") # Create a placeholder for serialized tensorflow.Example messages to be fed. serialized_tf_example = tf.placeholder(tf.string, name="tf_example") # Parse the tensorflow.Example looking for a feature named "x" with a single # floating point value. feature_configs = { "x": tf.FixedLenFeature([1], dtype=tf.float32), "x2": tf.FixedLenFeature([1], dtype=tf.float32, default_value=[0.0]) } tf_example = tf.parse_example(serialized_tf_example, feature_configs) # Use tf.identity() to assign name x = tf.identity(tf_example["x"], name="x") y = tf.add(tf.multiply(a, x), b, name="y") y2 = tf.add(tf.multiply(a, x), c, name="y2") x2 = tf.identity(tf_example["x2"], name="x2") y3 = tf.add(tf.multiply(a, x2), c, name="y3") # Create an assets file that can be saved and restored as part of the # SavedModel. original_assets_directory = "/tmp/original/export/assets" original_assets_filename = "foo.txt" original_assets_filepath = _write_assets(original_assets_directory, original_assets_filename) # Set up the assets collection. assets_filepath = tf.constant(original_assets_filepath) tf.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, assets_filepath) filename_tensor = tf.Variable(original_assets_filename, name="filename_tensor", trainable=False, collections=[]) assign_filename_op = filename_tensor.assign(original_assets_filename) # Set up the signature for Predict with input and output tensor # specification. predict_input_tensor = tf.saved_model.utils.build_tensor_info(x) predict_signature_inputs = {"x": predict_input_tensor} predict_output_tensor = tf.saved_model.utils.build_tensor_info(y) predict_signature_outputs = {"y": predict_output_tensor} predict_signature_def = ( tf.saved_model.signature_def_utils.build_signature_def( predict_signature_inputs, predict_signature_outputs, tf.saved_model.signature_constants.PREDICT_METHOD_NAME)) signature_def_map = { "regress_x_to_y": _build_regression_signature(serialized_tf_example, y), "regress_x_to_y2": _build_regression_signature(serialized_tf_example, y2), "regress_x2_to_y3": _build_regression_signature(x2, y3), "classify_x_to_y": _build_classification_signature(serialized_tf_example, y), tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: predict_signature_def } # Initialize all variables and then save the SavedModel. sess.run(tf.global_variables_initializer()) signature_def_map = { "regress_x_to_y": _build_regression_signature(serialized_tf_example, y), "regress_x_to_y2": _build_regression_signature(serialized_tf_example, y2), "regress_x2_to_y3": _build_regression_signature(x2, y3), "classify_x_to_y": _build_classification_signature(serialized_tf_example, y), "classify_x2_to_y3": _build_classification_signature(x2, y3), tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: predict_signature_def } if use_main_op: builder.add_meta_graph_and_variables( sess, [tf.saved_model.tag_constants.SERVING], signature_def_map=signature_def_map, assets_collection=tf.get_collection( tf.GraphKeys.ASSET_FILEPATHS), main_op=tf.group(main_op.main_op(), assign_filename_op)) else: builder.add_meta_graph_and_variables( sess, [tf.saved_model.tag_constants.SERVING], signature_def_map=signature_def_map, assets_collection=tf.get_collection( tf.GraphKeys.ASSET_FILEPATHS), legacy_init_op=tf.group(assign_filename_op)) builder.save(as_text)