diff --git a/README.md b/README.md index 1a701e9f..c7421f58 100644 --- a/README.md +++ b/README.md @@ -25,12 +25,10 @@ __0. Please make sure that you run in a Docker container or a virtual environmen The script pulls its own subset of TensorFlow, which might conflict with the existing TensorFlow/Keras installation. -__Note__: *Check that [`tf-nightly-2.0-preview`](https://pypi.org/project/tf-nightly-2.0-preview/#files) is available for your platform.* - -Most of the times, this means that you have to use Python 3.6.8 in your local -environment. To force Python 3.6.8 in your local project, you can install -[`pyenv`](https://github.com/pyenv/pyenv) and proceed as follows in the target -directory: +The converter supports both Python 2 and Python 3. But Python 3.6.8 is recommended +for your local environment. To force Python 3.6.8 in your local project, you can +install [`pyenv`](https://github.com/pyenv/pyenv) and proceed as follows in the +target directory: ```bash pyenv install 3.6.8 @@ -52,7 +50,35 @@ __1. Install the TensorFlow.js pip package:__ pip install tensorflowjs ``` -__2. Run the converter script provided by the pip package:__ +__2. Run the conversion script provided by the pip package:__ + +There are two way to trigger the model conversion: + +- The conversion wizard: tensorflowjs_wizard +- Regular conversion script: tensorflowjs_converter + +To start the conversion wizard: +```bash +tensorflowjs_wizard +``` + +This tool will walk you through the conversion process and provide you with +details explanations for each choice you need to make. Behind the scene it calls +the converter script (`tensorflowjs_converter`) in pip package. This is the easier +way to convert a single model. + +There is also dry run mode for the wizard, which will not perform the actual +conversion but only generate the command for `tensorflowjs_converter` command. +This command can be used in your own shell script. + +```bash +tensorflowjs_wizard --dryrun +``` + +To convert a batch of models or integrate the conversion process into your own +script, you should look into using the tensorflowjs_converter script. + +Here is detail information of parameters of the converter script. The converter expects a __TensorFlow SavedModel__, __TensorFlow Hub module__, __TensorFlow.js JSON__ format, __Keras HDF5 model__, or __tf.keras SavedModel__ diff --git a/python/requirements.txt b/python/requirements.txt index e898dea4..b3025df3 100644 --- a/python/requirements.txt +++ b/python/requirements.txt @@ -4,3 +4,4 @@ numpy==1.16.4 six==1.11.0 tensorflow==1.14.0 tensorflow-hub==0.5.0 +PyInquirer==1.0.3 diff --git a/python/setup.py b/python/setup.py index 976937a6..dae0342f 100644 --- a/python/setup.py +++ b/python/setup.py @@ -27,6 +27,7 @@ def _get_requirements(file): CONSOLE_SCRIPTS = [ 'tensorflowjs_converter = tensorflowjs.converters.converter:pip_main', + 'tensorflowjs_wizard = tensorflowjs.wizard:main', ] setuptools.setup( @@ -57,6 +58,7 @@ def _get_requirements(file): 'tensorflowjs.version', 'tensorflowjs.quantization', 'tensorflowjs.read_weights', + 'tensorflowjs.wizard', 'tensorflowjs.write_weights', 'tensorflowjs.converters', 'tensorflowjs.converters.common', diff --git a/python/tensorflowjs/BUILD b/python/tensorflowjs/BUILD index a278d735..b1d883ae 100644 --- a/python/tensorflowjs/BUILD +++ b/python/tensorflowjs/BUILD @@ -48,6 +48,13 @@ py_library( # `pip install tensorflow` or `pip install tensorflow-gpu`. ) +py_library( + name = "expect_PyInquirer_installed", + # This is a dummy rule used as a PyInquirer dependency in open-source. + # We expect PyInquirer to already be installed on the system, e.g. via + # `pip install PyInquirer`. +) + py_library( name = "quantization", srcs = ["quantization.py"], @@ -116,6 +123,29 @@ py_test( ], ) +py_test( + name = "wizard_test", + srcs = ["wizard_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":expect_numpy_installed", + ":wizard", + ], +) + +py_binary( + name = "wizard", + srcs = ["wizard.py"], + srcs_version = "PY2AND3", + deps = [ + ":converters/common", + ":converters/converter", + "//tensorflowjs:expect_h5py_installed", + "//tensorflowjs:expect_keras_installed", + "//tensorflowjs:expect_tensorflow_installed", + ], +) + # A filegroup BUILD target that includes all the op list json files in the # the op_list/ folder. The op_list folder itself is a symbolic link to the # actual op_list folder under src/. diff --git a/python/tensorflowjs/__init__.py b/python/tensorflowjs/__init__.py index 946f9495..d120fa33 100644 --- a/python/tensorflowjs/__init__.py +++ b/python/tensorflowjs/__init__.py @@ -21,5 +21,6 @@ from tensorflowjs import converters from tensorflowjs import quantization from tensorflowjs import version +from tensorflowjs import wizard __version__ = version.version diff --git a/python/tensorflowjs/converters/BUILD b/python/tensorflowjs/converters/BUILD index 03624e3b..463b2473 100644 --- a/python/tensorflowjs/converters/BUILD +++ b/python/tensorflowjs/converters/BUILD @@ -91,6 +91,7 @@ py_binary( srcs = ["converter.py"], srcs_version = "PY2AND3", deps = [ + ":common", ":keras_h5_conversion", ":keras_tfjs_loader", ":tf_saved_model_conversion_v2", diff --git a/python/tensorflowjs/converters/common.py b/python/tensorflowjs/converters/common.py index 25da4d64..f579ce15 100644 --- a/python/tensorflowjs/converters/common.py +++ b/python/tensorflowjs/converters/common.py @@ -30,6 +30,27 @@ GENERATED_BY_KEY = 'generatedBy' CONVERTED_BY_KEY = 'convertedBy' +# Model formats. +KERAS_SAVED_MODEL = 'keras_saved_model' +KERAS_MODEL = 'keras' +TF_SAVED_MODEL = 'tf_saved_model' +TF_HUB_MODEL = 'tf_hub' +TFJS_GRAPH_MODEL = 'tfjs_graph_model' +TFJS_LAYERS_MODEL = 'tfjs_layers_model' + +# CLI argument strings. +INPUT_PATH = 'input_path' +OUTPUT_PATH = 'output_path' +INPUT_FORMAT = 'input_format' +OUTPUT_FORMAT = 'output_format' +SIGNATURE_NAME = 'signature_name' +SAVED_MODEL_TAGS = 'saved_model_tags' +QUANTIZATION_BYTES = 'quantization_bytes' +SPLIT_WEIGHTS_BY_LAYER = 'split_weights_by_layer' +VERSION = 'version' +SKIP_OP_CHECK = 'skip_op_check' +STRIP_DEBUG_OPS = 'strip_debug_ops' +WEIGHT_SHARD_SIZE_BYTES = 'weight_shard_size_bytes' def get_converted_by(): """Get the convertedBy string for storage in model artifacts.""" diff --git a/python/tensorflowjs/converters/converter.py b/python/tensorflowjs/converters/converter.py index 4167e104..d7dd88ae 100644 --- a/python/tensorflowjs/converters/converter.py +++ b/python/tensorflowjs/converters/converter.py @@ -32,6 +32,7 @@ from tensorflowjs import quantization from tensorflowjs import version +from tensorflowjs.converters import common from tensorflowjs.converters import keras_h5_conversion as conversion from tensorflowjs.converters import keras_tfjs_loader from tensorflowjs.converters import tf_saved_model_conversion_v2 @@ -395,17 +396,17 @@ def _standardize_input_output_formats(input_format, output_format): 'Use --input_format=tfjs_layers_model instead.') input_format_is_keras = ( - input_format in ['keras', 'keras_saved_model']) + input_format in [common.KERAS_MODEL, common.KERAS_SAVED_MODEL]) input_format_is_tf = ( - input_format in ['tf_saved_model', 'tf_hub']) + input_format in [common.TF_SAVED_MODEL, common.TF_HUB_MODEL]) if output_format is None: # If no explicit output_format is provided, infer it from input format. if input_format_is_keras: - output_format = 'tfjs_layers_model' + output_format = common.TFJS_LAYERS_MODEL elif input_format_is_tf: - output_format = 'tfjs_graph_model' - elif input_format == 'tfjs_layers_model': - output_format = 'keras' + output_format = common.TFJS_GRAPH_MODEL + elif input_format == common.TFJS_LAYERS_MODEL: + output_format = common.KERAS_MODEL elif output_format == 'tensorflowjs': # https://github.com/tensorflow/tfjs/issues/1292: Remove the logic for the # explicit error message of the deprecated model type name 'tensorflowjs' @@ -434,12 +435,14 @@ def _parse_quantization_bytes(quantization_bytes): else: raise ValueError('Unsupported quantization bytes: %s' % quantization_bytes) - def get_arg_parser(): - """Create the argument parser for the converter binary.""" + """ + Create the argument parser for the converter binary. + """ + parser = argparse.ArgumentParser('TensorFlow.js model converters.') parser.add_argument( - 'input_path', + common.INPUT_PATH, nargs='?', type=str, help='Path to the input file or directory. For input format "keras", ' @@ -447,14 +450,18 @@ def get_arg_parser(): 'a SavedModel directory, session bundle directory, frozen model file, ' 'or TF-Hub module is expected.') parser.add_argument( - 'output_path', nargs='?', type=str, help='Path for all output artifacts.') + common.OUTPUT_PATH, + nargs='?', + type=str, + help='Path for all output artifacts.') parser.add_argument( - '--input_format', + '--%s' % common.INPUT_FORMAT, type=str, required=False, - default='tf_saved_model', - choices=set(['keras', 'keras_saved_model', - 'tf_saved_model', 'tf_hub', 'tfjs_layers_model', + default=common.TF_SAVED_MODEL, + choices=set([common.KERAS_MODEL, common.KERAS_SAVED_MODEL, + common.TF_SAVED_MODEL, common.TF_HUB_MODEL, + common.TFJS_LAYERS_MODEL, 'tensorflowjs']), help='Input format. ' 'For "keras", the input path can be one of the two following formats:\n' @@ -471,75 +478,64 @@ def get_arg_parser(): 'For "tf" formats, a SavedModel, frozen model, session bundle model, ' ' or TF-Hub module is expected.') parser.add_argument( - '--output_format', + '--%s' % common.OUTPUT_FORMAT, type=str, required=False, - choices=set(['keras', 'keras_saved_model', 'tfjs_layers_model', - 'tfjs_graph_model', 'tensorflowjs']), + choices=set([common.KERAS_MODEL, common.KERAS_SAVED_MODEL, + common.TFJS_LAYERS_MODEL, common.TFJS_GRAPH_MODEL, + 'tensorflowjs']), help='Output format. Default: tfjs_graph_model.') parser.add_argument( - '--signature_name', + '--%s' % common.SIGNATURE_NAME, type=str, default=None, help='Signature of the SavedModel Graph or TF-Hub module to load. ' 'Applicable only if input format is "tf_hub" or "tf_saved_model".') parser.add_argument( - '--saved_model_tags', + '--%s' % common.SAVED_MODEL_TAGS, type=str, default='serve', help='Tags of the MetaGraphDef to load, in comma separated string ' 'format. Defaults to "serve". Applicable only if input format is ' '"tf_saved_model".') parser.add_argument( - '--quantization_bytes', + '--%s' % common.QUANTIZATION_BYTES, type=int, choices=set(quantization.QUANTIZATION_BYTES_TO_DTYPES.keys()), help='How many bytes to optionally quantize/compress the weights to. 1- ' 'and 2-byte quantizaton is supported. The default (unquantized) size is ' '4 bytes.') parser.add_argument( - '--split_weights_by_layer', + '--%s' % common.SPLIT_WEIGHTS_BY_LAYER, action='store_true', help='Applicable to keras input_format only: Whether the weights from ' 'different layers are to be stored in separate weight groups, ' 'corresponding to separate binary weight files. Default: False.') parser.add_argument( - '--version', + '--%s' % common.VERSION, '-v', dest='show_version', action='store_true', help='Show versions of tensorflowjs and its dependencies') parser.add_argument( - '--skip_op_check', + '--%s' % common.SKIP_OP_CHECK, action='store_true', help='Skip op validation for TensorFlow model conversion.') parser.add_argument( - '--strip_debug_ops', + '--%s' % common.STRIP_DEBUG_OPS, type=bool, default=True, help='Strip debug ops (Print, Assert, CheckNumerics) from graph.') parser.add_argument( - '--weight_shard_size_bytes', + '--%s' % common.WEIGHT_SHARD_SIZE_BYTES, type=int, default=None, help='Shard size (in bytes) of the weight files. Currently applicable ' 'only to output_format=tfjs_layers_model.') return parser - -def pip_main(): - """Entry point for pip-packaged binary. - - Note that pip-packaged binary calls the entry method without - any arguments, which is why this method is needed in addition to the - `main` method below. - """ - main([' '.join(sys.argv[1:])]) - - -def main(argv): - args = get_arg_parser().parse_args(argv[0].split(' ')) - +def convert(arguments): + args = get_arg_parser().parse_args(arguments) if args.show_version: print('\ntensorflowjs %s\n' % version.version) print('Dependency versions:') @@ -556,9 +552,9 @@ def main(argv): weight_shard_size_bytes = 1024 * 1024 * 4 if args.weight_shard_size_bytes: - if args.output_format != 'tfjs_layers_model': + if args.output_format != common.TFJS_LAYERS_MODEL: raise ValueError( - 'The --weight_shard_size_byte flag is only supported under ' + 'The --weight_shard_size_bytes flag is only supported under ' 'output_format=tfjs_layers_model.') weight_shard_size_bytes = args.weight_shard_size_bytes @@ -575,7 +571,7 @@ def main(argv): if args.quantization_bytes else None) if (args.signature_name and input_format not in - ('tf_saved_model', 'tf_hub')): + (common.TF_SAVED_MODEL, common.TF_HUB_MODEL)): raise ValueError( 'The --signature_name flag is applicable only to "tf_saved_model" and ' '"tf_hub" input format, but the current input format is ' @@ -583,25 +579,27 @@ def main(argv): # TODO(cais, piyu): More conversion logics can be added as additional # branches below. - if input_format == 'keras' and output_format == 'tfjs_layers_model': + if (input_format == common.KERAS_MODEL and + output_format == common.TFJS_LAYERS_MODEL): dispatch_keras_h5_to_tfjs_layers_model_conversion( args.input_path, output_dir=args.output_path, quantization_dtype=quantization_dtype, split_weights_by_layer=args.split_weights_by_layer) - elif input_format == 'keras' and output_format == 'tfjs_graph_model': + elif (input_format == common.KERAS_MODEL and + output_format == common.TFJS_GRAPH_MODEL): dispatch_keras_h5_to_tfjs_graph_model_conversion( args.input_path, output_dir=args.output_path, quantization_dtype=quantization_dtype, skip_op_check=args.skip_op_check, strip_debug_ops=args.strip_debug_ops) - elif (input_format == 'keras_saved_model' and - output_format == 'tfjs_layers_model'): + elif (input_format == common.KERAS_SAVED_MODEL and + output_format == common.TFJS_LAYERS_MODEL): dispatch_keras_saved_model_to_tensorflowjs_conversion( args.input_path, args.output_path, quantization_dtype=quantization_dtype, split_weights_by_layer=args.split_weights_by_layer) - elif (input_format == 'tf_saved_model' and - output_format == 'tfjs_graph_model'): + elif (input_format == common.TF_SAVED_MODEL and + output_format == common.TFJS_GRAPH_MODEL): tf_saved_model_conversion_v2.convert_tf_saved_model( args.input_path, args.output_path, signature_def=args.signature_name, @@ -609,28 +607,28 @@ def main(argv): quantization_dtype=quantization_dtype, skip_op_check=args.skip_op_check, strip_debug_ops=args.strip_debug_ops) - elif (input_format == 'tf_hub' and - output_format == 'tfjs_graph_model'): + elif (input_format == common.TF_HUB_MODEL and + output_format == common.TFJS_GRAPH_MODEL): tf_saved_model_conversion_v2.convert_tf_hub_module( args.input_path, args.output_path, args.signature_name, args.saved_model_tags, skip_op_check=args.skip_op_check, strip_debug_ops=args.strip_debug_ops) - elif (input_format == 'tfjs_layers_model' and - output_format == 'keras'): + elif (input_format == common.TFJS_LAYERS_MODEL and + output_format == common.KERAS_MODEL): dispatch_tensorflowjs_to_keras_h5_conversion(args.input_path, args.output_path) - elif (input_format == 'tfjs_layers_model' and - output_format == 'keras_saved_model'): + elif (input_format == common.TFJS_LAYERS_MODEL and + output_format == common.KERAS_SAVED_MODEL): dispatch_tensorflowjs_to_keras_saved_model_conversion(args.input_path, args.output_path) - elif (input_format == 'tfjs_layers_model' and - output_format == 'tfjs_layers_model'): + elif (input_format == common.TFJS_LAYERS_MODEL and + output_format == common.TFJS_LAYERS_MODEL): dispatch_tensorflowjs_to_tensorflowjs_conversion( args.input_path, args.output_path, quantization_dtype=_parse_quantization_bytes(args.quantization_bytes), weight_shard_size_bytes=weight_shard_size_bytes) - elif (input_format == 'tfjs_layers_model' and - output_format == 'tfjs_graph_model'): + elif (input_format == common.TFJS_LAYERS_MODEL and + output_format == common.TFJS_GRAPH_MODEL): dispatch_tfjs_layers_model_to_tfjs_graph_conversion( args.input_path, args.output_path, quantization_dtype=_parse_quantization_bytes(args.quantization_bytes), @@ -641,6 +639,18 @@ def main(argv): 'Unsupported input_format - output_format pair: %s - %s' % (input_format, output_format)) +def pip_main(): + """Entry point for pip-packaged binary. + + Note that pip-packaged binary calls the entry method without + any arguments, which is why this method is needed in addition to the + `main` method below. + """ + main([' '.join(sys.argv[1:])]) + + +def main(argv): + convert(argv) if __name__ == '__main__': tf.app.run(main=main, argv=[' '.join(sys.argv[1:])]) diff --git a/python/tensorflowjs/wizard.py b/python/tensorflowjs/wizard.py new file mode 100644 index 00000000..d6e960d5 --- /dev/null +++ b/python/tensorflowjs/wizard.py @@ -0,0 +1,532 @@ +# Copyright 2019 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Interactive command line tool for tensorflow.js model conversion.""" + +from __future__ import print_function, unicode_literals + +import json +import os +import re +import sys + +import PyInquirer +import h5py +from tensorflow.core.framework import types_pb2 +from tensorflow.python.saved_model import loader_impl +from tensorflowjs.converters import converter +from tensorflowjs.converters import common + +# regex for recognizing valid url for TFHub module. +TFHUB_VALID_URL_REGEX = re.compile( + # http:// or https:// + r'^(http)s?://', re.IGNORECASE) + +# prompt style +prompt_style = PyInquirer.style_from_dict({ + PyInquirer.Token.Separator: '#6C6C6C', + PyInquirer.Token.QuestionMark: '#FF9D00 bold', + PyInquirer.Token.Selected: '#5F819D', + PyInquirer.Token.Pointer: '#FF9D00 bold', + PyInquirer.Token.Instruction: '', # default + PyInquirer.Token.Answer: '#5F819D bold', + PyInquirer.Token.Question: '', +}) + +def value_in_list(answers, key, values): + """Determine user's answer for the key is in the value list. + Args: + answer: Dict of user's answers to the questions. + key: question key. + values: List of values to check from. + """ + try: + value = answers[key] + return value in values + except KeyError: + return False + + +def get_tfjs_model_type(model_file): + with open(model_file) as f: + data = json.load(f) + return data['format'] + + +def detect_input_format(input_path): + """Determine the input format from model's input path or file. + Args: + input_path: string of the input model path + returns: + string: detected input format + string: normalized input path + """ + input_path = input_path.strip() + detected_input_format = None + if re.match(TFHUB_VALID_URL_REGEX, input_path): + detected_input_format = common.TF_HUB_MODEL + elif os.path.isdir(input_path): + if (any(fname.lower().endswith('saved_model.pb') + for fname in os.listdir(input_path))): + detected_input_format = common.TF_SAVED_MODEL + else: + for fname in os.listdir(input_path): + fname = fname.lower() + if fname.endswith('model.json'): + filename = os.path.join(input_path, fname) + if get_tfjs_model_type(filename) == common.TFJS_LAYERS_MODEL_FORMAT: + input_path = os.path.join(input_path, fname) + detected_input_format = common.TFJS_LAYERS_MODEL + break + elif os.path.isfile(input_path): + if h5py.is_hdf5(input_path): + detected_input_format = common.KERAS_MODEL + elif input_path.endswith('saved_model.pb'): + detected_input_format = common.TF_SAVED_MODEL + elif (input_path.endswith('model.json') and + get_tfjs_model_type(input_path) == common.TFJS_LAYERS_MODEL_FORMAT): + detected_input_format = common.TFJS_LAYERS_MODEL + + return detected_input_format, input_path + + +def input_path_message(answers): + """Determine question for model's input path. + Args: + answer: Dict of user's answers to the questions + """ + answer = answers[common.INPUT_FORMAT] + message = 'The original path seems to be wrong, ' + if answer == common.KERAS_MODEL: + return message + 'what is the path of input HDF5 file?' + elif answer == common.TF_HUB_MODEL: + return message + ("what is the TFHub module URL? \n" + "(i.e. https://tfhub.dev/google/imagenet/" + "mobilenet_v1_100_224/classification/1)") + else: + return message + 'what is the directory that contains the model?' + + +def validate_input_path(input_path, input_format): + """Validate the input path for given input format. + Args: + input_path: input path of the model. + input_format: model format string. + """ + path = os.path.expanduser(input_path.strip()) + if not path: + return 'Please enter a valid path' + if input_format == common.TF_HUB_MODEL: + if not re.match(TFHUB_VALID_URL_REGEX, path): + return """This is not an valid URL for TFHub module: %s, + We expect a URL that starts with http(s)://""" % path + elif not os.path.exists(path): + return 'Nonexistent path for the model: %s' % path + if input_format in (common.KERAS_SAVED_MODEL, common.TF_SAVED_MODEL): + is_dir = os.path.isdir(path) + if not is_dir and not path.endswith('saved_model.pb'): + return 'The path provided is not a directory or pb file: %s' % path + if (is_dir and + not any(f.endswith('saved_model.pb') for f in os.listdir(path))): + return 'Did not find a .pb file inside the directory: %s' % path + if input_format == common.TFJS_LAYERS_MODEL: + is_dir = os.path.isdir(path) + if not is_dir and not path.endswith('model.json'): + return 'The path provided is not a directory or json file: %s' % path + if is_dir and not any(f.endswith('model.json') for f in os.listdir(path)): + return 'Did not find the model.json file inside the directory: %s' % path + if input_format == common.KERAS_MODEL: + if not h5py.is_hdf5(path): + return 'The path provided is not a keras model file: %s' % path + return True + + +def expand_input_path(input_path): + """Expand the relative input path to absolute path, and add layers model file + name to the end if input format is `tfjs_layers_model`. + Args: + input_path: input path of the model. + Returns: + string: return expanded input path. + """ + input_path = os.path.expanduser(input_path.strip()) + is_dir = os.path.isdir(input_path) + if is_dir: + for fname in os.listdir(input_path): + if fname.endswith('.json'): + filename = os.path.join(input_path, fname) + return filename + return input_path + + +def output_path_exists(output_path): + """Check the existence of the output path. + Args: + output_path: input path of the model. + Returns: + bool: return true when the output directory exists. + """ + if os.path.exists(output_path): + return True + return False + + +def generate_arguments(params): + """Generate the tensorflowjs command string for the selected params. + Args: + params: user selected parameters for the conversion. + Returns: + list: the argument list for converter. + """ + args = [] + not_param_list = [common.INPUT_PATH, common.OUTPUT_PATH, + 'overwrite_output_path'] + no_false_param = [common.SPLIT_WEIGHTS_BY_LAYER, common.SKIP_OP_CHECK] + for key, value in sorted(params.items()): + if key not in not_param_list and value is not None: + if key in no_false_param: + if value is True: + args.append('--%s' % (key)) + else: + args.append('--%s=%s' % (key, value)) + + args.append(params[common.INPUT_PATH]) + args.append(params[common.OUTPUT_PATH]) + return args + + +def is_saved_model(input_format): + """Check if the input path contains saved model. + Args: + input_format: input model format. + Returns: + bool: whether this is for a saved model conversion. + """ + return input_format == common.TF_SAVED_MODEL or \ + input_format == common.KERAS_SAVED_MODEL + +def available_output_formats(answers): + """Generate the output formats for given input format. + Args: + ansowers: user selected parameter dict. + """ + input_format = answers[common.INPUT_FORMAT] + if input_format == common.KERAS_SAVED_MODEL: + return [{ + 'key': 'g', # shortcut key for the option + 'name': 'Tensorflow.js Graph Model', + 'value': common.TFJS_GRAPH_MODEL, + }, { + 'key': 'l', + 'name': 'TensoFlow.js Layers Model', + 'value': common.TFJS_LAYERS_MODEL, + }] + if input_format == common.TFJS_LAYERS_MODEL: + return [{ + 'key': 'k', + 'name': 'Keras Model (HDF5)', + 'value': common.KERAS_MODEL, + }, { + 'key': 'l', + 'name': 'TensoFlow.js Layers Model', + 'value': common.TFJS_LAYERS_MODEL, + }] + return [] + + +def available_tags(answers): + """Generate the available saved model tags from the proto file. + Args: + ansowers: user selected parameter dict. + """ + if is_saved_model(answers[common.INPUT_FORMAT]): + saved_model = loader_impl.parse_saved_model(answers[common.INPUT_PATH]) + tags = [] + for meta_graph in saved_model.meta_graphs: + tags.append(",".join(meta_graph.meta_info_def.tags)) + return tags + return [] + + +def available_signature_names(answers): + """Generate the available saved model signatures from the proto file + and selected tags. + Args: + ansowers: user selected parameter dict. + """ + if is_saved_model(answers[common.INPUT_FORMAT]): + path = answers[common.INPUT_PATH] + tags = answers[common.SAVED_MODEL_TAGS] + saved_model = loader_impl.parse_saved_model(path) + for meta_graph in saved_model.meta_graphs: + if tags == ",".join(meta_graph.meta_info_def.tags): + signatures = [] + for key in meta_graph.signature_def: + input_nodes = meta_graph.signature_def[key].inputs + output_nodes = meta_graph.signature_def[key].outputs + signatures.append( + {'value': key, + 'name': format_signature(key, input_nodes, output_nodes)}) + return signatures + return [] + + +def format_signature(name, input_nodes, output_nodes): + string = "signature name: %s\n" % name + string += " inputs: %s" % format_nodes(input_nodes) + string += " outputs: %s" % format_nodes(output_nodes) + return string + + +def format_nodes(nodes): + string = "%s of %s\n" % (3 if len(nodes) > 3 else len(nodes), len(nodes)) + count = 0 + for key in nodes: + value = nodes[key] + string += " name: %s, " % value.name + string += "dtype: %s, " % types_pb2.DataType.Name(value.dtype) + if value.tensor_shape.unknown_rank: + string += "shape: Unknown\n" + else: + string += "shape: %s\n" % [x.size for x in value.tensor_shape.dim] + count += 1 + if count >= 3: + break + return string + + +def input_format_string(base, target_format, detected_format): + if target_format == detected_format: + return base + ' *' + else: + return base + + +def input_format_message(detected_input_format): + message = 'What is your input model format? ' + if detected_input_format: + message += '(auto-detected format is marked with *)' + else: + message += '(model format cannot be detected.) ' + return message + + +def input_formats(detected_format): + formats = [{ + 'key': 'k', + 'name': input_format_string('Keras (HDF5)', common.KERAS_MODEL, + detected_format), + 'value': common.KERAS_MODEL + }, { + 'key': 'e', + 'name': input_format_string('Tensorflow Keras Saved Model', + common.KERAS_SAVED_MODEL, + detected_format), + 'value': common.KERAS_SAVED_MODEL, + }, { + 'key': 's', + 'name': input_format_string('Tensorflow Saved Model', + common.TF_SAVED_MODEL, + detected_format), + 'value': common.TF_SAVED_MODEL, + }, { + 'key': 'h', + 'name': input_format_string('TFHub Module', + common.TF_HUB_MODEL, + detected_format), + 'value': common.TF_HUB_MODEL, + }, { + 'key': 'l', + 'name': input_format_string('TensoFlow.js Layers Model', + common.TFJS_LAYERS_MODEL, + detected_format), + 'value': common.TFJS_LAYERS_MODEL, + }] + formats.sort(key=lambda x: x['value'] != detected_format) + return formats + + +def main(dryrun): + print('Welcome to TensorFlow.js Converter.') + input_path = [{ + 'type': 'input', + 'name': common.INPUT_PATH, + 'message': 'Please provide the path of model file or ' + 'the directory that contains model files. \n' + 'If you are converting TFHub module please provide the URL.', + 'filter': os.path.expanduser, + 'validate': + lambda path: 'Please enter a valid path' if not path else True + }] + + input_params = PyInquirer.prompt(input_path, style=prompt_style) + detected_input_format, normalized_path = detect_input_format( + input_params[common.INPUT_PATH]) + input_params[common.INPUT_PATH] = normalized_path + + formats = [ + { + 'type': 'list', + 'name': common.INPUT_FORMAT, + 'message': input_format_message(detected_input_format), + 'choices': input_formats(detected_input_format) + }, { + 'type': 'list', + 'name': common.OUTPUT_FORMAT, + 'message': 'What is your output format?', + 'choices': available_output_formats, + 'when': lambda answers: value_in_list(answers, common.INPUT_FORMAT, + (common.KERAS_SAVED_MODEL, + common.TFJS_LAYERS_MODEL)) + } + ] + format_params = PyInquirer.prompt(formats, input_params, style=prompt_style) + message = input_path_message(format_params) + + questions = [ + { + 'type': 'input', + 'name': common.INPUT_PATH, + 'message': message, + 'filter': expand_input_path, + 'validate': lambda value: validate_input_path( + value, format_params[common.INPUT_FORMAT]), + 'when': lambda answers: (not detected_input_format) + }, + { + 'type': 'list', + 'name': common.SAVED_MODEL_TAGS, + 'choices': available_tags, + 'message': 'What is tags for the saved model?', + 'when': lambda answers: (is_saved_model(answers[common.INPUT_FORMAT]) + and + (not common.OUTPUT_FORMAT in format_params + or format_params[common.OUTPUT_FORMAT] == + common.TFJS_GRAPH_MODEL)) + }, + { + 'type': 'list', + 'name': common.SIGNATURE_NAME, + 'message': 'What is signature name of the model?', + 'choices': available_signature_names, + 'when': lambda answers: (is_saved_model(answers[common.INPUT_FORMAT]) + and + (not common.OUTPUT_FORMAT in format_params + or format_params[common.OUTPUT_FORMAT] == + common.TFJS_GRAPH_MODEL)) + }, + { + 'type': 'list', + 'name': common.QUANTIZATION_BYTES, + 'message': 'Do you want to compress the model? ' + '(this will decrease the model precision.)', + 'choices': [{ + 'name': 'No compression, no accuracy loss.', + 'value': None + }, { + 'name': '2x compression, medium accuracy loss.', + 'value': 2 + }, { + 'name': '4x compression, highest accuracy loss.', + 'value': 1 + }] + }, + { + 'type': 'input', + 'name': common.WEIGHT_SHARD_SIZE_BYTES, + 'message': 'Please enter shard size (in bytes) of the weight files?', + 'default': str(4 * 1024 * 1024), + 'when': lambda answers: value_in_list(answers, common.OUTPUT_FORMAT, + (common.TFJS_LAYERS_MODEL)) + }, + { + 'type': 'confirm', + 'name': common.SPLIT_WEIGHTS_BY_LAYER, + 'message': 'Do you want to split weights by layers?', + 'default': False, + 'when': lambda answers: value_in_list(answers, common.INPUT_FORMAT, + (common.TFJS_LAYERS_MODEL)) + }, + { + 'type': 'confirm', + 'name': common.SKIP_OP_CHECK, + 'message': 'Do you want to skip op validation? \n' + 'This will allow conversion of unsupported ops, \n' + 'you can implement them as custom ops in tfjs-converter.', + 'default': False, + 'when': lambda answers: value_in_list(answers, common.INPUT_FORMAT, + (common.TF_SAVED_MODEL, + common.TF_HUB_MODEL)) + }, + { + 'type': 'confirm', + 'name': common.STRIP_DEBUG_OPS, + 'message': 'Do you want to strip debug ops? \n' + 'This will improve model execution performance.', + 'default': True, + 'when': lambda answers: value_in_list(answers, common.INPUT_FORMAT, + (common.TF_SAVED_MODEL, + common.TF_HUB_MODEL)) + } + ] + params = PyInquirer.prompt(questions, format_params, style=prompt_style) + + output_options = [ + { + 'type': 'input', + 'name': common.OUTPUT_PATH, + 'message': 'Which directory do you want to save ' + 'the converted model in?', + 'filter': lambda path: os.path.expanduser(path.strip()), + 'validate': lambda path: len(path) > 0 + }, { + 'type': 'confirm', + 'message': 'The output already directory exists, ' + 'do you want to overwrite it?', + 'name': 'overwrite_output_path', + 'default': False, + 'when': lambda ans: output_path_exists(ans[common.OUTPUT_PATH]) + } + ] + + while (not common.OUTPUT_PATH in params or + output_path_exists(params[common.OUTPUT_PATH]) and + not params['overwrite_output_path']): + params = PyInquirer.prompt(output_options, params, style=prompt_style) + + arguments = generate_arguments(params) + print('converter command generated:') + print('tensorflowjs_converter %s' % ' '.join(arguments)) + print('\n\n') + + if not dryrun: + converter.convert(arguments) + print('\n\nFile(s) generated by conversion:') + + print("Filename {0:25} Size(bytes)".format('')) + total_size = 0 + for basename in sorted(os.listdir(params[common.OUTPUT_PATH])): + filename = os.path.join(params[common.OUTPUT_PATH], basename) + size = os.path.getsize(filename) + print("{0:35} {1}".format(basename, size)) + total_size += size + print("Total size:{0:24} {1}".format('', total_size)) + + +if __name__ == '__main__': + if len(sys.argv) > 2 or len(sys.argv) == 2 and not sys.argv[1] == '--dryrun': + print("Usage: tensorflowjs_wizard [--dryrun]") + sys.exit(1) + dry_run = len(sys.argv) == 2 and sys.argv[1] == '--dryrun' + main(dry_run) diff --git a/python/tensorflowjs/wizard_test.py b/python/tensorflowjs/wizard_test.py new file mode 100644 index 00000000..bde7f14b --- /dev/null +++ b/python/tensorflowjs/wizard_test.py @@ -0,0 +1,233 @@ +# Copyright 2019 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import unittest +import tempfile +import json +import os +import shutil +import tensorflow as tf +from tensorflow import keras +from tensorflow.python.eager import def_function +from tensorflow.python.ops import variables +from tensorflow.python.training.tracking import tracking +from tensorflow.python.saved_model import save + +from tensorflowjs import wizard + +SAVED_MODEL_DIR = 'saved_model' +SAVED_MODEL_NAME = 'saved_model.pb' +HD5_FILE_NAME = 'test.h5' +LAYERS_MODEL_NAME = 'model.json' + + +class CliTest(unittest.TestCase): + def setUp(self): + self._tmp_dir = tempfile.mkdtemp() + super(CliTest, self).setUp() + + def tearDown(self): + if os.path.isdir(self._tmp_dir): + shutil.rmtree(self._tmp_dir) + super(CliTest, self).tearDown() + + def _create_layers_model(self): + data = {'format': 'layers-model'} + filename = os.path.join(self._tmp_dir, 'model.json') + with open(filename, 'a') as model_file: + json.dump(data, model_file) + + def _create_hd5_file(self): + input_tensor = keras.layers.Input((3,)) + dense1 = keras.layers.Dense( + 4, use_bias=True, kernel_initializer='ones', bias_initializer='zeros', + name='MyDense10')(input_tensor) + output = keras.layers.Dense( + 2, use_bias=False, kernel_initializer='ones', name='MyDense20')(dense1) + model = keras.models.Model(inputs=[input_tensor], outputs=[output]) + h5_path = os.path.join(self._tmp_dir, HD5_FILE_NAME) + print(h5_path) + model.save_weights(h5_path) + + def _create_saved_model(self): + """Test a basic model with functions to make sure functions are inlined.""" + input_data = tf.constant(1., shape=[1]) + root = tracking.AutoTrackable() + root.v1 = variables.Variable(3.) + root.v2 = variables.Variable(2.) + root.f = def_function.function(lambda x: root.v1 * root.v2 * x) + to_save = root.f.get_concrete_function(input_data) + + save_dir = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) + save.save(root, save_dir, to_save) + + def testOfValues(self): + answers = {'input_path': 'abc', 'input_format': '123'} + self.assertEqual(True, wizard.value_in_list(answers, 'input_path', ['abc'])) + self.assertEqual(False, wizard.value_in_list(answers, + 'input_path', ['abd'])) + self.assertEqual(False, wizard.value_in_list(answers, + 'input_format2', ['abc'])) + + def testInputPathMessage(self): + answers = {'input_format': 'keras'} + self.assertEqual("The original path seems to be wrong, " + "what is the path of input HDF5 file?", + wizard.input_path_message(answers)) + + answers = {'input_format': 'tf_hub'} + self.assertEqual("The original path seems to be wrong, " + "what is the TFHub module URL? \n" + "(i.e. https://tfhub.dev/google/imagenet/" + "mobilenet_v1_100_224/classification/1)", + wizard.input_path_message(answers)) + + answers = {'input_format': 'tf_saved_model'} + self.assertEqual("The original path seems to be wrong, " + "what is the directory that contains the model?", + wizard.input_path_message(answers)) + + def testValidateInputPathForTFHub(self): + self.assertNotEqual(True, + wizard.validate_input_path(self._tmp_dir, 'tf_hub')) + self.assertEqual(True, + wizard.validate_input_path("https://tfhub.dev/mobilenet", + 'tf_hub')) + + def testValidateInputPathForSavedModel(self): + self.assertNotEqual(True, wizard.validate_input_path( + self._tmp_dir, 'tf_saved_model')) + self._create_saved_model() + save_dir = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) + self.assertEqual(True, wizard.validate_input_path( + save_dir, 'tf_saved_model')) + + save_dir = os.path.join(self._tmp_dir, SAVED_MODEL_DIR, SAVED_MODEL_NAME) + self.assertEqual(True, wizard.validate_input_path( + save_dir, 'tf_saved_model')) + + def testValidateInputPathForKerasSavedModel(self): + self.assertNotEqual(True, wizard.validate_input_path( + self._tmp_dir, 'keras_saved_model')) + self._create_saved_model() + save_dir = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) + self.assertEqual(True, wizard.validate_input_path( + save_dir, 'keras_saved_model')) + + def testValidateInputPathForKerasModel(self): + self.assertNotEqual(True, + wizard.validate_input_path(self._tmp_dir, 'keras')) + self._create_hd5_file() + save_dir = os.path.join(self._tmp_dir, HD5_FILE_NAME) + self.assertEqual(True, wizard.validate_input_path( + save_dir, 'keras')) + + def testValidateInputPathForLayersModel(self): + self.assertNotEqual(True, + wizard.validate_input_path(self._tmp_dir, 'keras')) + self._create_layers_model() + save_dir = os.path.join(self._tmp_dir) + self.assertEqual(True, wizard.validate_input_path( + save_dir, 'tfjs_layers_model')) + + save_dir = os.path.join(self._tmp_dir, 'model.json') + self.assertEqual(True, wizard.validate_input_path( + save_dir, 'tfjs_layers_model')) + + def testOutputPathExist(self): + self.assertEqual(True, wizard.output_path_exists(self._tmp_dir)) + output_dir = os.path.join(self._tmp_dir, 'test') + self.assertNotEqual(True, wizard.output_path_exists(output_dir)) + + def testAvailableTags(self): + self._create_saved_model() + save_dir = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) + self.assertEqual(['serve'], wizard.available_tags( + {'input_path': save_dir, + 'input_format': 'tf_saved_model'})) + + def testAvailableSignatureNames(self): + self._create_saved_model() + save_dir = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) + self.assertEqual(['__saved_model_init_op', 'serving_default'], + [x['value'] for x in wizard.available_signature_names( + {'input_path': save_dir, + 'input_format': 'tf_saved_model', + 'saved_model_tags': 'serve'})]) + + def testGenerateCommandForSavedModel(self): + options = {'input_format': 'tf_saved_model', + 'input_path': 'tmp/saved_model', + 'saved_model_tags': 'test', + 'signature_name': 'test_default', + 'quantization_bytes': 2, + 'skip_op_check': False, + 'strip_debug_ops': True, + 'output_path': 'tmp/web_model'} + + self.assertEqual(['--input_format=tf_saved_model', + '--quantization_bytes=2', '--saved_model_tags=test', + '--signature_name=test_default', '--strip_debug_ops=True', + 'tmp/saved_model', 'tmp/web_model'], + wizard.generate_arguments(options)) + + def testGenerateCommandForKerasSavedModel(self): + options = {'input_format': 'tf_keras_saved_model', + 'output_format': 'tfjs_layers_model', + 'input_path': 'tmp/saved_model', + 'saved_model_tags': 'test', + 'signature_name': 'test_default', + 'quantization_bytes': 1, + 'skip_op_check': True, + 'strip_debug_ops': False, + 'output_path': 'tmp/web_model'} + + self.assertEqual(['--input_format=tf_keras_saved_model', + '--output_format=tfjs_layers_model', + '--quantization_bytes=1', '--saved_model_tags=test', + '--signature_name=test_default', '--skip_op_check', + '--strip_debug_ops=False', 'tmp/saved_model', + 'tmp/web_model'], + wizard.generate_arguments(options)) + + def testGenerateCommandForKerasModel(self): + options = {'input_format': 'keras', + 'input_path': 'tmp/model.HD5', + 'quantization_bytes': 1, + 'output_path': 'tmp/web_model'} + + self.assertEqual(['--input_format=keras', '--quantization_bytes=1', + 'tmp/model.HD5', 'tmp/web_model'], + wizard.generate_arguments(options)) + + def testGenerateCommandForLayerModel(self): + options = {'input_format': 'tfjs_layers_model', + 'output_format': 'keras', + 'input_path': 'tmp/model.json', + 'quantization_bytes': 1, + 'output_path': 'tmp/web_model'} + + self.assertEqual(['--input_format=tfjs_layers_model', + '--output_format=keras', + '--quantization_bytes=1', 'tmp/model.json', + 'tmp/web_model'], + wizard.generate_arguments(options)) + + +if __name__ == '__main__': + unittest.main()