Add sh_test to cover the converter binary (#412)

- Also in this change: Simplify the argument parsing code in converter.py
- Replace variable name "FLAGS" with "args" to remove the need for disabling pylint.

Fixes b/139308919

Author: caisq

python/setup.py

@@ -26,7 +26,7 @@ def _get_requirements(file):
         return requirements.readlines()
 
 CONSOLE_SCRIPTS = [
-    'tensorflowjs_converter = tensorflowjs.converters.converter:main',
+    'tensorflowjs_converter = tensorflowjs.converters.converter:pip_main',
 ]
 
 setuptools.setup(

python/tensorflowjs/converters/BUILD

@@ -101,6 +101,16 @@ py_binary(
     visibility = ["//visibility:public"],
 )
 
+py_binary(
+    name = "generate_test_model",
+    srcs = ["generate_test_model.py"],
+    testonly = True,
+    srcs_version = "PY2AND3",
+    deps = [
+        "//tensorflowjs:expect_tensorflow_installed",
+    ]
+)
+
 py_test(
     name = "converter_test",
     srcs = ["converter_test.py"],
@@ -113,3 +123,12 @@ py_test(
         "//tensorflowjs:version",
     ],
 )
+
+sh_test(
+    name = "converter_binary_test",
+    srcs = ["converter_binary_test.sh"],
+    data = [
+        ":converter",
+        ":generate_test_model",
+    ],
+)

python/tensorflowjs/converters/converter.py

@@ -527,47 +527,54 @@ def get_arg_parser():
   return parser
 
 
-def main():
-  try:
-    FLAGS
-  except NameError:
-    # This code path is added in addition to to the flags-parsing
-    # code in the `__name__ == '__main__'` branch below, because it is
-    # required by the pip-packaged binary case. The pip-packaged binary calls
-    # the `main()` method directly and therefore by passes the
-    # `__name__ == '__main__'` branch.
-    # pylint: disable=redefined-outer-name,invalid-name
-    FLAGS = get_arg_parser().parse_args()
-    # pylint: enable=redefined-outer-name,invalid-name
-
-  if FLAGS.show_version:
+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(' '))
+
+  if args.show_version:
     print('\ntensorflowjs %s\n' % version.version)
     print('Dependency versions:')
     print('  keras %s' % keras.__version__)
     print('  tensorflow %s' % tf.__version__)
     return
 
+  if not args.input_path:
+    raise ValueError(
+        'Missing input_path argument. For usage, use the --help flag.')
+  if not args.output_path:
+    raise ValueError(
+        'Missing output_path argument. For usage, use the --help flag.')
+
   weight_shard_size_bytes = 1024 * 1024 * 4
-  if FLAGS.weight_shard_size_bytes:
-    if  FLAGS.output_format != 'tfjs_layers_model':
+  if args.weight_shard_size_bytes:
+    if  args.output_format != 'tfjs_layers_model':
       raise ValueError(
           'The --weight_shard_size_byte flag is only supported under '
           'output_format=tfjs_layers_model.')
-    weight_shard_size_bytes = FLAGS.weight_shard_size_bytes
+    weight_shard_size_bytes = args.weight_shard_size_bytes
 
-  if FLAGS.input_path is None:
+  if args.input_path is None:
     raise ValueError(
         'Error: The input_path argument must be set. '
         'Run with --help flag for usage information.')
 
   input_format, output_format = _standardize_input_output_formats(
-      FLAGS.input_format, FLAGS.output_format)
+      args.input_format, args.output_format)
 
   quantization_dtype = (
-      quantization.QUANTIZATION_BYTES_TO_DTYPES[FLAGS.quantization_bytes]
-      if FLAGS.quantization_bytes else None)
+      quantization.QUANTIZATION_BYTES_TO_DTYPES[args.quantization_bytes]
+      if args.quantization_bytes else None)
 
-  if (FLAGS.signature_name and input_format not in
+  if (args.signature_name and input_format not in
       ('tf_saved_model', 'tf_hub')):
     raise ValueError(
         'The --signature_name flag is applicable only to "tf_saved_model" and '
@@ -578,65 +585,62 @@ def main():
   #   branches below.
   if input_format == 'keras' and output_format == 'tfjs_layers_model':
     dispatch_keras_h5_to_tfjs_layers_model_conversion(
-        FLAGS.input_path, output_dir=FLAGS.output_path,
+        args.input_path, output_dir=args.output_path,
         quantization_dtype=quantization_dtype,
-        split_weights_by_layer=FLAGS.split_weights_by_layer)
+        split_weights_by_layer=args.split_weights_by_layer)
   elif input_format == 'keras' and output_format == 'tfjs_graph_model':
     dispatch_keras_h5_to_tfjs_graph_model_conversion(
-        FLAGS.input_path, output_dir=FLAGS.output_path,
+        args.input_path, output_dir=args.output_path,
         quantization_dtype=quantization_dtype,
-        skip_op_check=FLAGS.skip_op_check,
-        strip_debug_ops=FLAGS.strip_debug_ops)
+        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'):
     dispatch_keras_saved_model_to_tensorflowjs_conversion(
-        FLAGS.input_path, FLAGS.output_path,
+        args.input_path, args.output_path,
         quantization_dtype=quantization_dtype,
-        split_weights_by_layer=FLAGS.split_weights_by_layer)
+        split_weights_by_layer=args.split_weights_by_layer)
   elif (input_format == 'tf_saved_model' and
         output_format == 'tfjs_graph_model'):
     tf_saved_model_conversion_v2.convert_tf_saved_model(
-        FLAGS.input_path, FLAGS.output_path,
-        signature_def=FLAGS.signature_name,
-        saved_model_tags=FLAGS.saved_model_tags,
+        args.input_path, args.output_path,
+        signature_def=args.signature_name,
+        saved_model_tags=args.saved_model_tags,
         quantization_dtype=quantization_dtype,
-        skip_op_check=FLAGS.skip_op_check,
-        strip_debug_ops=FLAGS.strip_debug_ops)
+        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'):
     tf_saved_model_conversion_v2.convert_tf_hub_module(
-        FLAGS.input_path, FLAGS.output_path, FLAGS.signature_name,
-        FLAGS.saved_model_tags, skip_op_check=FLAGS.skip_op_check,
-        strip_debug_ops=FLAGS.strip_debug_ops)
+        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'):
-    dispatch_tensorflowjs_to_keras_h5_conversion(FLAGS.input_path,
-                                                 FLAGS.output_path)
+    dispatch_tensorflowjs_to_keras_h5_conversion(args.input_path,
+                                                 args.output_path)
   elif (input_format == 'tfjs_layers_model' and
         output_format == 'keras_saved_model'):
-    dispatch_tensorflowjs_to_keras_saved_model_conversion(FLAGS.input_path,
-                                                          FLAGS.output_path)
+    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'):
     dispatch_tensorflowjs_to_tensorflowjs_conversion(
-        FLAGS.input_path, FLAGS.output_path,
-        quantization_dtype=_parse_quantization_bytes(FLAGS.quantization_bytes),
+        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'):
     dispatch_tfjs_layers_model_to_tfjs_graph_conversion(
-        FLAGS.input_path, FLAGS.output_path,
-        quantization_dtype=_parse_quantization_bytes(FLAGS.quantization_bytes),
-        skip_op_check=FLAGS.skip_op_check,
-        strip_debug_ops=FLAGS.strip_debug_ops)
+        args.input_path, args.output_path,
+        quantization_dtype=_parse_quantization_bytes(args.quantization_bytes),
+        skip_op_check=args.skip_op_check,
+        strip_debug_ops=args.strip_debug_ops)
   else:
     raise ValueError(
         'Unsupported input_format - output_format pair: %s - %s' %
         (input_format, output_format))
 
 
 if __name__ == '__main__':
-  # pylint: disable=redefined-outer-name,invalid-name
-  FLAGS, unparsed = get_arg_parser().parse_known_args()
-  # pylint: enable=redefined-outer-name,invalid-name
-  tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
+  tf.app.run(main=main, argv=[' '.join(sys.argv[1:])])

python/tensorflowjs/converters/converter_binary_test.sh

@@ -0,0 +1,61 @@
+# 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.
+# ==============================================================================
+
+set -e
+
+GENERATE_BIN="${TEST_SRCDIR}/org_tensorflow_js/tensorflowjs/converters/generate_test_model"
+CONVERTER_BIN="${TEST_SRCDIR}/org_tensorflow_js/tensorflowjs/converters/converter"
+
+# 1. Test tf_saved_model --> tfjs_graph_model conversion.
+SAVED_MODEL_DIR="$(mktemp -d)"
+echo "Genearting TF SavedModel for testing..."
+"${GENERATE_BIN}" "${SAVED_MODEL_DIR}" --model_type tf_saved_model
+echo "Done genearting TF SavedModel for testing at ${SAVED_MODEL_DIR}"
+
+OUTPUT_DIR="${SAVED_MODEL_DIR}_converted"
+"${CONVERTER_BIN}" \
+    --input_format tf_saved_model \
+    --output_format tfjs_graph_model \
+    "${SAVED_MODEL_DIR}" \
+    "${OUTPUT_DIR}"
+
+if [[ ! -d "${OUTPUT_DIR}" ]]; then
+  echo "ERROR: Failed to find conversion output directory: ${OUTPUT_DIR}" 1>&2
+  exit 1
+fi
+
+# Clean up files.
+rm -rf "${SAVED_MODEL_DIR}" "${OUTPUT_DIR}"
+
+# 2. Test keras HDF5 --> tfjs_layers_model conversion.
+KERAS_H5_PATH="$(mktemp).h5"
+echo "Genearting Keras HDF5 model for testing..."
+"${GENERATE_BIN}" "${KERAS_H5_PATH}" --model_type tf_keras_h5
+echo "Done genearting Keras HDF5 model for testing at ${KERAS_H5_PATH}"
+
+OUTPUT_H5_PATH="${KERAS_H5_PATH}_converted.h5"
+"${CONVERTER_BIN}" \
+    --input_format keras \
+    --output_format tfjs_layers_model \
+    "${KERAS_H5_PATH}" \
+    "${OUTPUT_H5_PATH}"
+
+if [[ ! -d "${OUTPUT_H5_PATH}" ]]; then
+  echo "ERROR: Failed to find conversion output directory: ${OUTPUT_H5_PATH}" 1>&2
+  exit 1
+fi
+
+# Clean up files.
+rm -rf "${KERAS_H5_PATH}" "${OUTPUT_H5_PATH}"

python/tensorflowjs/converters/generate_test_model.py

@@ -0,0 +1,68 @@
+# 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.
+# ==============================================================================
+"""A binary that generates saved model artifacts for testing."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import os
+import sys
+
+import tensorflow as tf
+
+tf.enable_eager_execution()
+
+
+def parse_args():
+  parser = argparse.ArgumentParser(
+      'Generates saved model artifacts for testing.')
+  parser.add_argument(
+      'output_path',
+      type=str,
+      help='Model output path.')
+  parser.add_argument(
+      '--model_type',
+      type=str,
+      required=True,
+      choices=set(['tf_keras_h5', 'tf_saved_model']),
+      help='Model format to generate.')
+  return parser.parse_known_args()
+
+
+def main(_):
+
+  if args.model_type == 'tf_keras_h5':
+    model = tf.keras.Sequential()
+    model.add(tf.keras.layers.Dense(5, activation='relu', input_shape=(8,)))
+    model.add(tf.keras.layers.Dense(1, activation='sigmoid'))
+    model.save(os.path.join(args.output_path))
+  elif args.model_type == 'tf_saved_model':
+    class TimesThreePlusOne(tf.train.Checkpoint):
+
+      @tf.function(input_signature=[
+          tf.TensorSpec(shape=None, dtype=tf.float32)])
+      def compute(self, x):
+        return x * 3.0 + 1.0
+
+    tf.saved_model.save(TimesThreePlusOne(), args.output_path)
+  else:
+    raise ValueError('Unrecognized model type: %s' % args.model_type)
+
+
+if __name__ == '__main__':
+  args, unparsed = parse_args()
+  tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)