Add TFJS conversion (#1614)

Add TFJS conversion

Author: TomWildenhain-Microsoft

ci_build/azure_pipelines/templates/job_generator.yml

@@ -12,6 +12,7 @@ parameters:
   report_coverage: 'False'
   artifact_name: 'single_test_coverage'
   skip_tflite_tests: 'True'
+  skip_tfjs_tests: 'True'
   skip_tf_tests: 'False'
 
 jobs:
@@ -29,7 +30,7 @@ jobs:
                 ${{ each onnx_backend in parameters.onnx_backends }}:
                   ${{ each onnx_backend_version in onnx_backend.value }}:
                     ${{ each onnx_opset in parameters.onnx_opsets }}:
-                      ${{ format('{0} python{1}{2} tf{3} onnx{4} {5}{6}{7}', platform, python_version, replace(replace(parameters.skip_tflite_tests,'True', ''), 'False', ' tflite'), tf_version, onnx_version, replace(format('opset{0} ', onnx_opset), 'opset ', ''), onnx_backend.key, onnx_backend_version) }}:
+                      ${{ format('{0} python{1}{2}{3} tf{4} onnx{5} {6}{7}{8}', platform, python_version, replace(replace(parameters.skip_tflite_tests,'True', ''), 'False', ' tflite'), replace(replace(parameters.skip_tfjs_tests,'True', ''), 'False', ' tfjs'), tf_version, onnx_version, replace(format('opset{0} ', onnx_opset), 'opset ', ''), onnx_backend.key, onnx_backend_version) }}:
                         ${{ if eq(platform, 'linux') }}:
                           CI_VM_IMAGE: 'ubuntu-16.04'
                         ${{ if eq(platform, 'windows') }}:
@@ -46,6 +47,7 @@ jobs:
                         CI_ONNX_BACKEND_VERSION: '${{ onnx_backend_version }}'
                         CI_SKIP_TF_TESTS: '${{ parameters.skip_tf_tests }}'
                         CI_SKIP_TFLITE_TESTS: '${{ parameters.skip_tflite_tests }}'
+                        CI_SKIP_TFJS_TESTS: '${{ parameters.skip_tfjs_tests }}'
 
                         ${{ if eq(tf_version, '') }}:
                           CI_PIP_TF_NAME: 'tensorflow'

ci_build/azure_pipelines/templates/setup.yml

@@ -25,6 +25,12 @@ steps:
       pip install --index-url https://test.pypi.org/simple/ ort-nightly
     fi
 
+    if [[ $CI_SKIP_TFJS_TESTS == "False" ]] ;
+    then
+      pip install tensorflowjs
+      npm install @tensorflow/tfjs
+    fi
+
     if [[ $CI_TF_VERSION == 2.* ]] ;
     then
       pip install onnxruntime-extensions

ci_build/azure_pipelines/templates/unit_test.yml

@@ -3,6 +3,7 @@
 parameters:
     onnx_opsets: ['14', '13', '12', '11', '10', '9', '8']
     skip_tflite_tests: 'True'
+    skip_tfjs_tests: 'True'
     skip_tf_tests: 'False'
 
 steps:
@@ -11,6 +12,7 @@ steps:
       export TF2ONNX_TEST_BACKEND=$CI_ONNX_BACKEND
       export TF2ONNX_TEST_OPSET=$CI_ONNX_OPSET
       export TF2ONNX_SKIP_TFLITE_TESTS=$CI_SKIP_TFLITE_TESTS
+      export TF2ONNX_SKIP_TFJS_TESTS=$CI_SKIP_TFJS_TESTS
       export TF2ONNX_SKIP_TF_TESTS=$CI_SKIP_TF_TESTS
       python -m pytest --cov=tf2onnx --cov-report=term --disable-pytest-warnings -r s tests --cov-append
     timeoutInMinutes: 15
@@ -19,4 +21,5 @@ steps:
     env:
         CI_ONNX_OPSET: '${{ onnx_opset }}'
         CI_SKIP_TFLITE_TESTS: '${{ parameters.skip_tflite_tests }}'
+        CI_SKIP_TFJS_TESTS: '${{ parameters.skip_tfjs_tests }}'
         CI_SKIP_TF_TESTS: '${{ parameters.skip_tf_tests }}'

ci_build/azure_pipelines/unit_test.yml

@@ -3,6 +3,19 @@
 stages:
   - stage:
     jobs:
+    - template: 'templates/job_generator.yml'
+      parameters:
+        # TFJS tf 2.5
+        python_versions: ['3.9']
+        tf_versions: ['2.5.0']
+        onnx_opsets: ['']
+        skip_tfjs_tests: 'False'
+        skip_tf_tests: 'True'
+        job:
+          steps:
+          - template: 'unit_test.yml'
+        report_coverage: 'True'
+
     - template: 'templates/job_generator.yml'
       parameters:
         # TFLite tf 2.5

tests/backend_test_base.py

@@ -18,6 +18,7 @@
 from tensorflow.python.ops import lookup_ops
 import onnx
 from common import get_test_config
+from tfjs_runner import run_tfjs
 from tf2onnx import utils
 from tf2onnx.tfonnx import process_tf_graph
 from tf2onnx import optimizer
@@ -113,6 +114,8 @@ def assert_results_equal(self, expected, actual, rtol, atol, mtol=None,
                     decode = np.vectorize(lambda x: x.replace(b'\x00', b'').decode('UTF-8'))
                     expected_val_str = decode(expected_val)
                     self.assertAllEqual(expected_val_str, actual_val)
+                elif expected_val.dtype.kind == 'U':
+                    self.assertAllEqual(expected_val, actual_val)
                 else:
                     if mtol is not None:
                         expected_val = np.minimum(expected_val, mtol)
@@ -189,11 +192,21 @@ def freeze_and_run_tf(self, func, feed_dict, outputs, as_session, premade_placeh
                 tf.import_graph_def(graph_def, name='')
                 graph_def = tf_optimize(list(feed_dict.keys()), outputs, graph_def, fold_constant=constant_fold)
 
-        model_path = os.path.join(self.test_data_directory, self._testMethodName + "_after_tf_optimize.pb")
-        utils.save_protobuf(model_path, graph_def)
-        self.logger.debug("created file  %s", model_path)
         return result, graph_def, initialized_tables
 
+    def convert_to_tfjs(self, graph_def_path, output_names):
+        from tensorflowjs.converters import converter
+        tfjs_path = os.path.join(self.test_data_directory, self._testMethodName + "_tfjs")
+        try:
+            converter.convert([graph_def_path, tfjs_path, '--input_format', 'tf_frozen_model',
+                               '--output_node_names', ','.join(output_names)])
+        except ValueError:
+            return None
+        model_path = os.path.join(tfjs_path, 'model.json')
+        if not os.path.exists(model_path):
+            return None
+        return model_path
+
     def convert_to_tflite(self, graph_def, feed_dict, outputs):
         if not feed_dict:
             return None   # Can't make TFlite model with no inputs
@@ -306,6 +319,7 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
                       use_custom_ops=False):
         test_tf = not self.config.skip_tf_tests
         test_tflite = not self.config.skip_tflite_tests
+        test_tfjs = not self.config.skip_tfjs_tests
         run_tfl_consistency_test = test_tf and test_tflite and self.config.run_tfl_consistency_test
         # optional - passed to process_tf_graph
         if process_args is None:
@@ -323,6 +337,15 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
             self.freeze_and_run_tf(func, feed_dict, output_names_with_port, as_session,
                                    premade_placeholders, large_model, constant_fold)
 
+        graph_def_path = os.path.join(self.test_data_directory, self._testMethodName + "_after_tf_optimize.pb")
+        utils.save_protobuf(graph_def_path, graph_def)
+        self.logger.debug("created file  %s", graph_def_path)
+
+        if test_tfjs:
+            tfjs_path = self.convert_to_tfjs(graph_def_path, output_names_with_port)
+            if tfjs_path is None:
+                test_tfjs = False
+
         if test_tflite:
             tflite_path = self.convert_to_tflite(graph_def, feed_dict, output_names_with_port)
             test_tflite = tflite_path is not None and self.tflite_has_supported_types(tflite_path)
@@ -383,8 +406,39 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
             if graph_validator:
                 self.assertTrue(graph_validator(g))
 
+        if test_tfjs:
+            try:
+                tfjs_res = run_tfjs(tfjs_path, feed_dict, self.test_data_directory)
+            except RuntimeError as e:
+                ignored_errors = ["is not yet supported", "Operands could not be broadcast together",
+                                  "unknown dtype null", "must be [NaN", "Cannot read property 'name' of undefined",
+                                  "Either strides or dilations must be 1", "does not support"]
+                if any(err in str(e) for err in ignored_errors):
+                    test_tfjs = False
+                else:
+                    raise e
+
+        if test_tfjs:
+            g = process_tf_graph(None, opset=self.config.opset,
+                                 input_names=list(feed_dict.keys()),
+                                 output_names=None,
+                                 target=self.config.target,
+                                 tfjs_path=tfjs_path,
+                                 **process_args)
+            g = optimizer.optimize_graph(g)
+            onnx_tfjs_res = self.run_backend(g, None, onnx_feed_dict, large_model,
+                                             postfix="_from_tfjs", use_custom_ops=use_custom_ops)
+
+            self.assert_results_equal(tfjs_res, onnx_tfjs_res, rtol, atol, mtol, check_value, check_shape,
+                                      check_dtype=False)
+            self.assert_shapes_correct(g, self.config.allow_missing_shapes, not self.config.skip_onnx_checker)
+
+            if graph_validator:
+                self.assertTrue(graph_validator(g))
+
+
         if g is None:
-            raise unittest.SkipTest("Both tf and tflite marked to skip")
+            raise unittest.SkipTest("tf, tflite, and tfjs marked to skip")
         return g
 
     def save_onnx_model(self, model_proto, feed_dict, postfix="", external_tensor_storage=None):

tests/common.py

@@ -31,6 +31,7 @@
     "check_opset_max_version",
     "skip_tf2",
     "skip_tflite",
+    "skip_tfjs",
     "requires_tflite",
     "check_opset_after_tf_version",
     "check_target",
@@ -58,6 +59,7 @@ def __init__(self):
         self.target = os.environ.get("TF2ONNX_TEST_TARGET", ",".join(constants.DEFAULT_TARGET)).split(',')
         self.backend = os.environ.get("TF2ONNX_TEST_BACKEND", "onnxruntime")
         self.skip_tflite_tests = os.environ.get("TF2ONNX_SKIP_TFLITE_TESTS", "FALSE").upper() == "TRUE"
+        self.skip_tfjs_tests = os.environ.get("TF2ONNX_SKIP_TFJS_TESTS", "FALSE").upper() == "TRUE"
         self.skip_tf_tests = os.environ.get("TF2ONNX_SKIP_TF_TESTS", "FALSE").upper() == "TRUE"
         self.skip_onnx_checker = False
         self.allow_missing_shapes = False
@@ -102,6 +104,7 @@ def __str__(self):
                             "opset={}".format(self.opset),
                             "target={}".format(self.target),
                             "skip_tflite_tests={}".format(self.skip_tflite_tests),
+                            "skip_tfjs_tests={}".format(self.skip_tfjs_tests),
                             "skip_tf_tests={}".format(self.skip_tf_tests),
                             "run_tfl_consistency_test={}".format(self.run_tfl_consistency_test),
                             "backend={}".format(self.backend),
@@ -181,12 +184,31 @@ def skip_tf2(message=""):
     return unittest.skipIf(tf_loader.is_tf2(), reason)
 
 
+def skip_tfjs(message=""):
+    """ Skip the tfjs conversion for this test """
+    config = get_test_config()
+    reason = _append_message("test disabled for tfjs", message)
+    if config.skip_tf_tests and config.skip_tflite_tests:
+        # If we are skipping tf and tflite also, there is no reason to run this test
+        return unittest.skip(reason)
+    def decorator(func):
+        def test(self):
+            tmp = config.skip_tfjs_tests
+            config.skip_tfjs_tests = True
+            try:
+                func(self)
+            finally:
+                config.skip_tfjs_tests = tmp
+        return test
+    return decorator
+
+
 def skip_tflite(message=""):
     """ Skip the tflite conversion for this test """
     config = get_test_config()
     reason = _append_message("test disabled for tflite", message)
-    if config.skip_tf_tests:
-        # If we are skipping tf also, there is no reason to run this test
+    if config.skip_tf_tests and config.skip_tfjs_tests:
+        # If we are skipping tf and tfjs also, there is no reason to run this test
         return unittest.skip(reason)
     def decorator(func):
         def test(self):

tests/run_tfjs.js

@@ -0,0 +1,143 @@
+// SPDX-License-Identifier: Apache-2.0
+
+
+/* Helper script to run tfjs models. Uses custom .json representation to encode tensors for inputs/outputs
+ * Example usage:
+ * 
+ *     node run_tfjs.js mymodel/model.json input.json output.json
+ */
+
+const tf = require('@tensorflow/tfjs');
+
+const fs = require('fs');
+const http = require('http');
+const path = require('path');
+const { exit } = require('process');
+
+const [, , modelPath, inputPath, outputPath] = process.argv;
+
+// Hide tfjs first use message complaining about lack of GPU
+tf.backend().firstUse = false;
+
+if (process.argv[2] == '--test') {
+    // dtype = 'float32'|'int32'|'bool'|'complex64'|'string'
+
+    const floatTensor = tf.tensor([1.1, 2.2, 3.3, 4.4], [2, 2], 'float32');
+    const intTensor = tf.tensor([1, 2, 3, 4], [2, 2], 'int32');
+    const boolTensor = tf.tensor([true, false, true, true], [2, 2], 'bool');
+    const complexTensor = tf.complex([1.1, 2.2, 3.3, 4.4], [10., 20., 30., 40.]).reshape([2, 2]);
+    const stringTensor = tf.tensor(['Hello world', '♦♥♠♣', '', 'Tensors'], [2, 2], 'string');
+
+    const tensors = [floatTensor, intTensor, boolTensor, complexTensor, stringTensor];
+    tensors.forEach(function (tensor) {
+        const tensorEnc = tensorToJson(tensor);
+        const tensorDec = tensorFromJson(tensorEnc);
+        if (tensor.toString() != tensorDec.toString()) {
+            console.log("Tensor:")
+            tensor.print()
+            console.log("Decoded tensor:")
+            tensorDec.print()
+            throw "Test failure"
+        }
+    });
+
+    console.log("All tests pass.")
+    exit(0)
+}
+
+const modelDir = path.dirname(modelPath);
+const modelName = path.basename(modelPath);
+
+// tf.loadGraphModel expects a url not a local file, so we serve it on localhost 
+http.createServer(function (req, res) {
+    fs.readFile(modelDir + req.url, function (err, data) {
+        if (err) {
+            res.writeHead(404);
+            res.end(JSON.stringify(err));
+            return;
+        }
+        res.writeHead(200);
+        res.end(data);
+    });
+}).listen(8080);
+
+function tensorToJson(tensor) {
+    if (tensor.dtype != 'string') {
+        const data = tensor.dataSync()
+        const byteArray = new Uint8Array(data.buffer, data.byteOffset, data.byteLength)
+        const dataEnc = Buffer.from(byteArray).toString('base64');
+        return {
+            dataEnc: dataEnc,
+            shape: tensor.shape,
+            dtype: tensor.dtype
+        };
+    }
+
+    return {
+        data: tensor.dataSync(),
+        shape: tensor.shape,
+        dtype: tensor.dtype
+    };
+}
+
+function getTypedArrayConstructor(dtype) {
+    if (dtype == 'complex64') {
+        return Float32Array;
+    }
+    return tf.util.getTypedArrayFromDType(dtype).constructor;
+}
+
+function tensorFromJson(json) {
+    let data = json.data;
+    if (data == undefined) {
+        const arrayType = getTypedArrayConstructor(json.dtype);
+        data = new arrayType(new Uint8Array(Buffer.from(json.dataEnc, 'base64')).buffer);
+    }
+    if (json.dtype == 'complex64') {
+        const floatTensor = tf.tensor(data, [data.length], 'float32').reshape([-1, 2]).transpose();
+        return tf.complex(floatTensor.gather(0), floatTensor.gather(1)).reshape(json.shape);
+    }
+    return tf.tensor(data, json.shape, json.dtype);
+}
+
+function inputFromJson(json) {
+    // Input can be a tensor, list of tensors, or mapping of input names to tensors.
+    if (Array.isArray(json)) {
+        return json.map(tensorFromJson);
+    }
+    const keys = Object.keys(json);
+    if (keys.length == 0 || json[keys[0]].dtype != undefined) {
+        const result = {};
+        keys.forEach(k => { result[k] = tensorFromJson(json[k]); });
+        return result;
+    }
+    return tensorFromJson(json);
+}
+
+function outputToJson(out) {
+    // Output can be a tensor, list of tensors, or mapping of output names to tensors.
+    if (Array.isArray(out)) {
+        return out.map(tensorToJson);
+    }
+    if (out instanceof tf.Tensor) {
+        return tensorToJson(out);
+    }
+    const result = {};
+    Object.keys(out).forEach(k => { result[k] = tensorToJson(out[k]); });
+    return result;
+}
+
+async function main() {
+    const model = await tf.loadGraphModel('http://localhost:8080/' + modelName);
+    const inputString = fs.readFileSync(inputPath, 'utf8');
+    const inputJson = JSON.parse(inputString);
+    const input = inputFromJson(inputJson);
+
+    const output = await model.executeAsync(input);
+
+    const outputJson = outputToJson(output);
+    const outputString = JSON.stringify(outputJson);
+    fs.writeFileSync(outputPath, outputString, 'utf8');
+}
+
+main().then(() => exit(0)).catch((err) => { console.error(err); exit(1) })