diff --git a/tfjs-converter/python/requirements-exp.txt b/tfjs-converter/python/requirements-exp.txt new file mode 100644 index 00000000000..3d3772bf904 --- /dev/null +++ b/tfjs-converter/python/requirements-exp.txt @@ -0,0 +1,8 @@ +h5py>=2.8.0 +numpy>=1.16.4,<1.19.0 +six>=1.12.0 +tf-nightly-cpu>=2.4.0.dev20200806,<3 +tensorflow-hub==0.7.0 +PyInquirer==1.0.3 +pylint==1.9.4; python_version < '3.0' +pylint==2.5.0; python_version > '3.0' diff --git a/tfjs-converter/python/run-python-tests.sh b/tfjs-converter/python/run-python-tests.sh index 7d1fd39c5c9..209a74a046d 100755 --- a/tfjs-converter/python/run-python-tests.sh +++ b/tfjs-converter/python/run-python-tests.sh @@ -16,6 +16,12 @@ # A script that runs all Python unit tests in tfjs-layers. +function print_usage() { + echo "Usage:" + echo " run-python-tests.sh " + echo +} + set -e SCRIPTS_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" @@ -28,8 +34,26 @@ TMP_VENV_DIR="$(mktemp -u).venv" virtualenv -p "python" "${TMP_VENV_DIR}" source "${TMP_VENV_DIR}/bin/activate" -pip install -r "${SCRIPTS_DIR}/requirements-dev.txt" +# There is one argument (requirements_file), please update this constant when +# you adding more arguments. +ARGS_COUNT=1 + +# Default requirements file name. +REQ_FILE="${SCRIPTS_DIR}/requirements-dev.txt" + +# Show the usage message if there are too many arguments. +if [[ $# > ARGS_COUNT ]]; then + print_usage + exit 1 +fi + +# Use the user specified requirements file name. +if [[ $# == 1 ]]; then + REQ_FILE=$1 +fi +pip install -r "${REQ_FILE}" +# Run pylint for tensorflowjs directory cd "${SCRIPTS_DIR}" pylint --rcfile=.pylintrc tensorflowjs @@ -45,5 +69,6 @@ echo echo "All tests passed." echo +# Clean up deactivate rm -rf "${TMP_VENV_DIR}" diff --git a/tfjs-converter/python/tensorflowjs/converters/BUILD b/tfjs-converter/python/tensorflowjs/converters/BUILD index 7e00c6db7ab..56ab69f16f7 100644 --- a/tfjs-converter/python/tensorflowjs/converters/BUILD +++ b/tfjs-converter/python/tensorflowjs/converters/BUILD @@ -159,6 +159,7 @@ py_library( "//tensorflowjs:expect_numpy_installed", "//tensorflowjs:expect_tensorflow_hub_installed", "//tensorflowjs:expect_tensorflow_installed", + "//tensorflowjs:expect_graph_transforms_installed", "//tensorflowjs:resource_loader", "//tensorflowjs:version", "//tensorflowjs:write_weights", diff --git a/tfjs-converter/python/tensorflowjs/converters/common.py b/tfjs-converter/python/tensorflowjs/converters/common.py index 2775d5aedb6..2f76019d8b5 100644 --- a/tfjs-converter/python/tensorflowjs/converters/common.py +++ b/tfjs-converter/python/tensorflowjs/converters/common.py @@ -59,6 +59,7 @@ STRIP_DEBUG_OPS = 'strip_debug_ops' WEIGHT_SHARD_SIZE_BYTES = 'weight_shard_size_bytes' CONTROL_FLOW_V2 = 'control_flow_v2' +EXPERIMENTS = 'experiments' def get_converted_by(): """Get the convertedBy string for storage in model artifacts.""" diff --git a/tfjs-converter/python/tensorflowjs/converters/converter.py b/tfjs-converter/python/tensorflowjs/converters/converter.py index ed682ac1b73..96631be5402 100644 --- a/tfjs-converter/python/tensorflowjs/converters/converter.py +++ b/tfjs-converter/python/tensorflowjs/converters/converter.py @@ -105,7 +105,8 @@ def dispatch_keras_h5_to_tfjs_graph_model_conversion( skip_op_check=False, strip_debug_ops=False, weight_shard_size_bytes=1024 * 1024 * 4, - control_flow_v2=False): + control_flow_v2=False, + experiments=False): """ Convert a keras HDF5-format model to tfjs GraphModel artifacts. @@ -120,6 +121,8 @@ def dispatch_keras_h5_to_tfjs_graph_model_conversion( strip_debug_ops: Bool whether to allow unsupported debug ops. weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. + control_flow_v2: Bool whether to enable control flow v2 ops. + experiments: Bool enable experimental features. """ if not os.path.exists(h5_path): @@ -143,7 +146,8 @@ def dispatch_keras_h5_to_tfjs_graph_model_conversion( skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, weight_shard_size_bytes=weight_shard_size_bytes, - control_flow_v2=control_flow_v2) + control_flow_v2=control_flow_v2, + experiments=experiments) # Clean up the temporary SavedModel directory. shutil.rmtree(temp_savedmodel_dir) @@ -331,7 +335,9 @@ def dispatch_tfjs_layers_model_to_tfjs_graph_conversion( quantization_dtype_map=None, skip_op_check=False, strip_debug_ops=False, - weight_shard_size_bytes=1024 * 1024 * 4): + weight_shard_size_bytes=1024 * 1024 * 4, + control_flow_v2=False, + experiments=False): """Converts a TensorFlow.js Layers Model to TensorFlow.js Graph Model. This conversion often benefits speed of inference, due to the graph @@ -348,7 +354,8 @@ def dispatch_tfjs_layers_model_to_tfjs_graph_conversion( strip_debug_ops: Bool whether to allow unsupported debug ops. weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. - + control_flow_v2: Bool whether to enable control flow v2 ops. + experiments: Bool enable experimental features. Raises: ValueError, if `config_json_path` is not a path to a valid JSON file, or if h5_path points to an existing directory. @@ -382,7 +389,9 @@ def dispatch_tfjs_layers_model_to_tfjs_graph_conversion( quantization_dtype_map=quantization_dtype_map, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, - weight_shard_size_bytes=weight_shard_size_bytes) + weight_shard_size_bytes=weight_shard_size_bytes, + control_flow_v2=control_flow_v2, + experiments=experiments) # Clean up temporary HDF5 file. os.remove(temp_h5_path) @@ -575,9 +584,16 @@ def get_arg_parser(): '"tf_frozen_model".') parser.add_argument( '--%s' % common.CONTROL_FLOW_V2, - type=str, + type=bool, + default=False, help='Enable control flow v2 ops, this would improve inference ' 'performance on models with branches or loops.') + parser.add_argument( + '--%s' % common.EXPERIMENTS, + type=bool, + default=False, + help='Enable experimental features, you should only enable this flag ' + 'when using Python3 and TensorFlow nightly build.') return parser def convert(arguments): @@ -660,7 +676,8 @@ def convert(arguments): skip_op_check=args.skip_op_check, strip_debug_ops=args.strip_debug_ops, weight_shard_size_bytes=weight_shard_size_bytes, - control_flow_v2=args.control_flow_v2) + control_flow_v2=args.control_flow_v2, + experiments=args.experiments) elif (input_format == common.KERAS_SAVED_MODEL and output_format == common.TFJS_LAYERS_MODEL): dispatch_keras_saved_model_to_tensorflowjs_conversion( @@ -678,7 +695,8 @@ def convert(arguments): skip_op_check=args.skip_op_check, strip_debug_ops=args.strip_debug_ops, weight_shard_size_bytes=weight_shard_size_bytes, - control_flow_v2=args.control_flow_v2) + control_flow_v2=args.control_flow_v2, + experiments=args.experiments) elif (input_format == common.TF_HUB_MODEL and output_format == common.TFJS_GRAPH_MODEL): tf_saved_model_conversion_v2.convert_tf_hub_module( @@ -689,7 +707,8 @@ def convert(arguments): skip_op_check=args.skip_op_check, strip_debug_ops=args.strip_debug_ops, weight_shard_size_bytes=weight_shard_size_bytes, - control_flow_v2=args.control_flow_v2) + control_flow_v2=args.control_flow_v2, + experiments=args.experiments) elif (input_format == common.TFJS_LAYERS_MODEL and output_format == common.KERAS_MODEL): dispatch_tensorflowjs_to_keras_h5_conversion(args.input_path, @@ -711,7 +730,9 @@ def convert(arguments): quantization_dtype_map=quantization_dtype_map, skip_op_check=args.skip_op_check, strip_debug_ops=args.strip_debug_ops, - weight_shard_size_bytes=weight_shard_size_bytes) + weight_shard_size_bytes=weight_shard_size_bytes, + control_flow_v2=args.control_flow_v2, + experiments=args.experiments) elif (input_format == common.TF_FROZEN_MODEL and output_format == common.TFJS_GRAPH_MODEL): tf_saved_model_conversion_v2.convert_tf_frozen_model( @@ -719,7 +740,8 @@ def convert(arguments): quantization_dtype_map=quantization_dtype_map, skip_op_check=args.skip_op_check, strip_debug_ops=args.strip_debug_ops, - weight_shard_size_bytes=weight_shard_size_bytes) + weight_shard_size_bytes=weight_shard_size_bytes, + experiments=args.experiments) else: raise ValueError( 'Unsupported input_format - output_format pair: %s - %s' % diff --git a/tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py b/tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py index 832d26197e9..c126a97b9f4 100644 --- a/tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py +++ b/tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py @@ -112,7 +112,7 @@ def _run_grappler(config, graph_def, graph, signature_def): def optimize_graph(graph, signature_def, output_graph, tf_version, quantization_dtype_map=None, skip_op_check=False, strip_debug_ops=False, - weight_shard_size_bytes=1024 * 1024 * 4): + weight_shard_size_bytes=1024 * 1024 * 4, experiments=False): """Takes a Python Graph object and optimizes the graph. Args: @@ -149,6 +149,9 @@ def optimize_graph(graph, signature_def, output_graph, 'pruning', 'constfold', 'arithmetic', 'dependency', 'pruning', 'constfold', 'arithmetic', 'dependency' ] + if experiments: + rewriter_config.experimental_disable_compressed_tensor_optimization = True + if strip_debug_ops: rewriter_config.optimizers.insert(0, 'debug_stripper') @@ -391,7 +394,8 @@ def convert_tf_frozen_model(frozen_model_path, quantization_dtype_map=None, skip_op_check=False, strip_debug_ops=False, - weight_shard_size_bytes=1024 * 1024 * 4): + weight_shard_size_bytes=1024 * 1024 * 4, + experiments=False): """Convert frozen model and check the model compatibility with Tensorflow.js. Optimize and convert the model to Tensorflow.js format, when the model passes the compatiblity check. @@ -409,6 +413,7 @@ def convert_tf_frozen_model(frozen_model_path, strip_debug_ops: Bool whether to strip debug ops. weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. + experiments: Bool enable experimental features. """ if not os.path.exists(output_dir): @@ -424,7 +429,8 @@ def convert_tf_frozen_model(frozen_model_path, quantization_dtype_map=quantization_dtype_map, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, - weight_shard_size_bytes=weight_shard_size_bytes) + weight_shard_size_bytes=weight_shard_size_bytes, + experiments=experiments) def convert_tf_saved_model(saved_model_dir, output_dir, signature_def='serving_default', @@ -433,7 +439,7 @@ def convert_tf_saved_model(saved_model_dir, skip_op_check=False, strip_debug_ops=False, weight_shard_size_bytes=1024 * 1024 * 4, - control_flow_v2=False): + control_flow_v2=False, experiments=False): """Freeze the SavedModel and check the model compatibility with Tensorflow.js. Optimize and convert the model to Tensorflow.js format, when the model passes @@ -457,6 +463,7 @@ def convert_tf_saved_model(saved_model_dir, weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. control_flow_v2: Bool whether to enable control flow v2 ops. + experiments: Bool enable experimental features. """ if signature_def is None: signature_def = 'serving_default' @@ -493,12 +500,73 @@ def convert_tf_saved_model(saved_model_dir, signature = _build_signature_def( frozen_graph, inputs, concrete_func.outputs) + # Check if the TransformGraph is available to be imported, this package is + # available in g3 but not in oss version of TensorFlow. + transform_graph_available = True + try: + from tensorflow.tools.graph_transforms import TransformGraph # pylint: disable=C0415 + except: # pylint: disable=W0702 + transform_graph_available = False + + # Define the strip graph functions when TransformGraph is available, this will + # strip the unused nodes from the graph. + if transform_graph_available: + def _gen_strip_unused_nodes_transformation(tensor): + """Generate `strip_unused_nodes()` transformation for the + input `tensor`. + """ + name = tensor.name.split(":")[0] + # "" -> "int64" + dtype = repr(tensor.dtype).replace("tf.", "") + # TensorShape([1, 128]) -> "1,128" + shape = ",".join(str(dim) for dim in tensor.shape) + return ('name="{}", type_for_name={}, ' + 'shape_for_name="{}"').format(name, dtype, shape) + + def _strip_unused_nodes(frozen_graph, concrete_func, output_node_names): + # Find the names of the input nodes needed to extract the minimal + # inference graph. This is particularly useful for cases when the concrete + # function contains nodes that do not contribute the inference computation + # defined by the input/output pair. This would also eliminate op + # unsupported error caused by nodes outside of the minial infrerence + # graph. + input_node_names = [] + for input_tensor in concrete_func.inputs: + if not input_tensor.dtype == 'resource': + op_name = input_tensor.name.split(':')[0] + # The graph freezing may turn the original inputs into constants, or + # remove them from the graph, so we need to ignore those. + try: + op = frozen_graph.get_operation_by_name(op_name) + if op.type != 'Const': + input_node_names.append(op_name) + except KeyError: + # The original input was removed when the graph was frozen. + continue + + graph_transformations = [ + "strip_unused_nodes(" + ", ".join( + _gen_strip_unused_nodes_transformation(tensor) + for tensor in concrete_func.inputs + if not tensor.dtype == 'resource') + ")" + ] + stripped_graph_def = TransformGraph( + frozen_graph.as_graph_def(), input_node_names, output_node_names, + graph_transformations) + with tf.Graph().as_default() as stripped_graph: + tf.import_graph_def(stripped_graph_def, name='') + return stripped_graph + + frozen_graph = _strip_unused_nodes( + frozen_graph, concrete_func, output_node_names) + optimize_graph(frozen_graph, signature, output_graph, model.tensorflow_version, quantization_dtype_map=quantization_dtype_map, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, - weight_shard_size_bytes=weight_shard_size_bytes) + weight_shard_size_bytes=weight_shard_size_bytes, + experiments=experiments) def load_and_initialize_hub_module(module_path, signature='default'): """Loads graph of a TF-Hub module and initializes it into a session. @@ -549,7 +617,8 @@ def load_and_initialize_hub_module(module_path, signature='default'): def convert_tf_hub_module_v1(module_path, output_dir, signature='default', quantization_dtype_map=None, skip_op_check=False, strip_debug_ops=False, - weight_shard_size_bytes=1024 * 1024 * 4): + weight_shard_size_bytes=1024 * 1024 * 4, + experiments=False): """Freeze the TF-Hub module and check compatibility with Tensorflow.js. Optimize and convert the TF-Hub module to Tensorflow.js format, if it passes @@ -569,6 +638,7 @@ def convert_tf_hub_module_v1(module_path, output_dir, strip_debug_ops: Bool whether to strip debug ops. weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. + experiments: Bool enable experimental features. """ if signature is None: @@ -609,7 +679,8 @@ def convert_tf_hub_module_v1(module_path, output_dir, quantization_dtype_map=quantization_dtype_map, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, - weight_shard_size_bytes=weight_shard_size_bytes) + weight_shard_size_bytes=weight_shard_size_bytes, + experiments=experiments) finally: # Clean up the temp files. if os.path.exists(frozen_file): @@ -621,7 +692,7 @@ def convert_tf_hub_module(module_handle, output_dir, quantization_dtype_map=None, skip_op_check=False, strip_debug_ops=False, weight_shard_size_bytes=1024 * 1024 * 4, - control_flow_v2=False): + control_flow_v2=False, experiments=False): """Conversion for TF Hub modules V1 and V2. See convert_tf_hub_module and convert_tf_saved_model. @@ -642,6 +713,7 @@ def convert_tf_hub_module(module_handle, output_dir, weight_shard_size_bytes: Shard size (in bytes) of the weight files. The size of each weight file will be <= this value. control_flow_v2: Bool whether to enable control flow v2 ops. + experiments: Bool enable experimental features. """ module_path = hub.resolve(module_handle) # TODO(vbardiovskyg): We can remove this v1 code path once loading of all v1 @@ -652,7 +724,8 @@ def convert_tf_hub_module(module_handle, output_dir, convert_tf_hub_module_v1(module_path, output_dir, signature, quantization_dtype_map, skip_op_check, strip_debug_ops, - weight_shard_size_bytes) + weight_shard_size_bytes, + experiments=experiments) else: print("Loading the module using TF 2.X interface from %s." % module_path) if signature is None: @@ -665,4 +738,5 @@ def convert_tf_hub_module(module_handle, output_dir, skip_op_check=skip_op_check, strip_debug_ops=strip_debug_ops, weight_shard_size_bytes=weight_shard_size_bytes, - control_flow_v2=control_flow_v2) + control_flow_v2=control_flow_v2, + experiments=experiments)