def testTrainedMnistSavedModel(self): """Test mnist SavedModel, trained with dummy data and small steps.""" # Build classifier classifier = estimator.Estimator( model_fn=model_fn, params={ "data_format": "channels_last" # tflite format }) # Train and pred for serving classifier.train(input_fn=dummy_input_fn, steps=2) image = array_ops.placeholder(dtypes.float32, [None, 28, 28]) pred_input_fn = estimator.export.build_raw_serving_input_receiver_fn({ "image": image, }) # Export SavedModel saved_model_dir = os.path.join(self.get_temp_dir(), "mnist_savedmodel") classifier.export_savedmodel(saved_model_dir, pred_input_fn) # Convert to tflite and test output saved_model_name = os.listdir(saved_model_dir)[0] saved_model_final_dir = os.path.join(saved_model_dir, saved_model_name) # TODO(zhixianyan): no need to limit output_arrays to `Softmax' # once b/74205001 fixed and argmax implemented in tflite. result = convert_saved_model.freeze_saved_model( saved_model_dir=saved_model_final_dir, input_arrays=None, input_shapes=None, output_arrays=["Softmax"], tag_set=set([tag_constants.SERVING]), signature_key=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY) self.assertTrue(result)
def testTrainedMnistSavedModel(self): """Test mnist SavedModel, trained with dummy data and small steps.""" # Build classifier classifier = estimator.Estimator( model_fn=model_fn, params={ "data_format": "channels_last" # tflite format }) # Train and pred for serving classifier.train(input_fn=dummy_input_fn, steps=2) image = array_ops.placeholder(dtypes.float32, [None, 28, 28]) pred_input_fn = estimator.export.build_raw_serving_input_receiver_fn({ "image": image, }) # Export SavedModel saved_model_dir = os.path.join(self.get_temp_dir(), "mnist_savedmodel") classifier.export_savedmodel(saved_model_dir, pred_input_fn) # Convert to tflite and test output saved_model_name = os.listdir(saved_model_dir)[0] saved_model_final_dir = os.path.join(saved_model_dir, saved_model_name) # TODO(zhixianyan): no need to limit output_arrays to `Softmax' # once b/74205001 fixed and argmax implemented in tflite. result = convert_saved_model.freeze_saved_model( saved_model_dir=saved_model_final_dir, input_arrays=None, input_shapes=None, output_arrays=["Softmax"], tag_set=None, signature_key=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY) self.assertTrue(result)
def from_saved_model( cls, saved_model_dir, input_arrays=None, input_shapes=None, output_arrays=None, tag_set=None, signature_key=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY): """Creates a TocoConverter class from a SavedModel. Args: saved_model_dir: SavedModel directory to convert. input_arrays: List of input tensors to freeze graph with. Uses input arrays from SignatureDef when none are provided. (default None) input_shapes: Map of strings representing input tensor names to list of integers representing input shapes (e.g., {"foo": : [1, 16, 16, 3]}). Automatically determined when input shapes is None (e.g., {"foo" : None}). (default None) output_arrays: List of output tensors to freeze graph with. Uses output arrays from SignatureDef when none are provided. (default None) tag_set: Set of tags identifying the MetaGraphDef within the SavedModel to analyze. All tags in the tag set must be present. (default "serve") signature_key: Key identifying SignatureDef containing inputs and outputs. Returns: TocoConverter class. """ if tag_set is None: tag_set = set([tag_constants.SERVING]) result = freeze_saved_model(saved_model_dir, input_arrays, input_shapes, output_arrays, tag_set, signature_key) return cls( graph_def=result[0], input_tensors=result[1], output_tensors=result[2])
def _convertSavedModel(self, saved_model_dir, input_arrays=None, input_shapes=None, output_arrays=None, tag_set=None, signature_key=None): graph_def, in_tensors, out_tensors = convert_saved_model.freeze_saved_model( saved_model_dir=saved_model_dir, input_arrays=input_arrays, input_shapes=input_shapes, output_arrays=output_arrays, tag_set=tag_set, signature_key=signature_key) return graph_def, in_tensors, out_tensors
def _convertSavedModel(self, saved_model_dir, input_arrays=None, input_shapes=None, output_arrays=None, tag_set=None, signature_key=None): if tag_set is None: tag_set = set([tag_constants.SERVING]) if signature_key is None: signature_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY graph_def, in_tensors, out_tensors = convert_saved_model.freeze_saved_model( saved_model_dir=saved_model_dir, input_arrays=input_arrays, input_shapes=input_shapes, output_arrays=output_arrays, tag_set=tag_set, signature_key=signature_key) return graph_def, in_tensors, out_tensors
def from_saved_model(cls, saved_model_dir, input_arrays=None, input_shapes=None, output_arrays=None, tag_set=None, signature_key=None): """Creates a TocoConverter class from a SavedModel. Args: saved_model_dir: SavedModel directory to convert. input_arrays: List of input tensors to freeze graph with. Uses input arrays from SignatureDef when none are provided. (default None) input_shapes: Dict of strings representing input tensor names to list of integers representing input shapes (e.g., {"foo" : [1, 16, 16, 3]}). Automatically determined when input shapes is None (e.g., {"foo" : None}). (default None) output_arrays: List of output tensors to freeze graph with. Uses output arrays from SignatureDef when none are provided. (default None) tag_set: Set of tags identifying the MetaGraphDef within the SavedModel to analyze. All tags in the tag set must be present. (default set("serve")) signature_key: Key identifying SignatureDef containing inputs and outputs. (default DEFAULT_SERVING_SIGNATURE_DEF_KEY) Returns: TocoConverter class. """ if tag_set is None: tag_set = set([tag_constants.SERVING]) if signature_key is None: signature_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY result = freeze_saved_model(saved_model_dir, input_arrays, input_shapes, output_arrays, tag_set, signature_key) return cls(graph_def=result[0], input_tensors=result[1], output_tensors=result[2])