Add conversion of _FusedDepthwiseConv2dNative and Prelu for TFJS (#1628)

* Add conversion of _FusedDepthwiseConv2dNative

Signed-off-by: Tom Wildenhain <[email protected]>

* Implement PRelu for tfjs

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/backend_test_base.py

@@ -420,7 +420,7 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
 
         if test_tfjs:
             try:
-                tfjs_res = run_tfjs(tfjs_path, feed_dict, self.test_data_directory)
+                tfjs_res = run_tfjs(tfjs_path, feed_dict)
             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",

tests/run_pretrained_models.py

@@ -451,7 +451,7 @@ def run_tflite():
                 inputs[k] = self.make_input(v)
             if not self.skip_tensorflow:
                 logger.info("Running TFJS")
-                tf_results = run_tfjs(tfjs_path, inputs, dir_name)
+                tf_results = run_tfjs(tfjs_path, inputs, outputs)
                 logger.info("TFJS OK")
 
         if not self.run_tf_frozen:

tests/run_pretrained_models.yaml

@@ -636,3 +636,37 @@ posenet_mobilenet_quantized_2_075_tfjs:
   rtol: 0.1
   ptol: 0.2
   atol: 0.005
+
+blazeposedetector_tfjs:
+  tf_min_version: 2.1
+  disabled: false
+  url: https://tfhub.dev/mediapipe/tfjs-model/blazeposedetector/1/default/1?tfjs-format=compressed
+  model: "model.json"
+  opset_constraints:
+    "onnx":
+      "min": 10
+  model_type: tfjs
+  input_get: get_beach
+  #force_input_shape: True  # ORT doesn't implement autopadding for convs with dilations
+  inputs:
+    "input:0": [1, 224, 224, 3]
+  outputs:
+    - Identity:0
+  rtol: 0.05
+  atol: 0.0005
+
+facemesh_tfjs:
+  tf_min_version: 2.1
+  disabled: false
+  url: https://tfhub.dev/mediapipe/tfjs-model/facemesh/1/default/1?tfjs-format=compressed
+  model: "model.json"
+  model_type: tfjs
+  input_get: get_beach
+  inputs:
+    "input_1:0": [1, 192, 192, 3]
+  outputs:
+    - Identity:0
+    - Identity_1:0
+    - Identity_2:0
+  rtol: 0.05
+  atol: 0.0005

tests/run_tfjs.js

@@ -15,7 +15,7 @@ const http = require('http');
 const path = require('path');
 const { exit } = require('process');
 
-const [, , modelPath, inputPath, outputPath] = process.argv;
+const [, , modelPath, inputPath, outputPath, ...args] = process.argv;
 
 // Hide tfjs first use message complaining about lack of GPU
 tf.backend().firstUse = false;
@@ -147,8 +147,11 @@ async function main() {
     const inputString = fs.readFileSync(inputPath, 'utf8');
     const inputJson = JSON.parse(inputString);
     const input = inputFromJson(inputJson);
-
-    const output = await model.executeAsync(input);
+    let outputs = null;
+    if (args && args[0] == '--outputs') {
+        outputs = args[1].split(',');
+    }
+    const output = await model.executeAsync(input, outputs);
 
     const outputJson = outputToJson(output);
     const outputString = JSON.stringify(outputJson);

tests/tfjs_runner.py

@@ -67,7 +67,7 @@ def json_to_output(obj):
     return [json_to_numpy(obj)]
 
 
-def run_tfjs(tfjs_path, inputs, working_dir):
+def run_tfjs(tfjs_path, inputs, outputs=None):
     """
     Given the path to the model.json of a tfjs model, a dict mapping input names to numpy arrays, and a working
     directory, runs the model on the inputs and returns the resulting arrays or raises a RuntimeException. Calls
@@ -79,11 +79,15 @@ def run_tfjs(tfjs_path, inputs, working_dir):
     output_path = os.path.join(working_dir, 'output.json')
     stderr_path = os.path.join(working_dir, 'stderr.txt')
 
+    command = ['node', script_path, tfjs_path, input_path, output_path]
+    if outputs is not None:
+        command.extend(['--outputs', ','.join(outputs)])
+
     with open(input_path, 'wt') as f:
         json.dump(inputs_to_json(inputs), f)
 
     with open(stderr_path, 'wb') as f:
-        result = subprocess.run(['node', script_path, tfjs_path, input_path, output_path], stderr=f, check=False)
+        result = subprocess.run(command, stderr=f, check=False)
     if result.returncode != 0:
         with open(stderr_path, 'rt') as f:
             err = f.read()

tf2onnx/onnx_opset/math.py

@@ -70,6 +70,13 @@ def version_9(cls, ctx, node, **kwargs):
         pass
 
 
+@tf_op(["Prelu"], onnx_op="PRelu")
+class Prelu:
+    @classmethod
+    def version_1(cls, ctx, node, **kwargs):
+        pass
+
+
 def make_min_or_max_op(ctx, op_type, inputs, outputs,
                        output_shapes=None, output_dtypes=None):
     # support more dtype

tf2onnx/rewriter/fused_op_rewriter.py

@@ -11,7 +11,7 @@
 
 def rewrite_fused_ops(g, ops):
     for node in ops:
-        if node.type in ["_FusedConv2D", "_FusedMatMul"]:
+        if node.type in ["_FusedConv2D", "_FusedMatMul", "_FusedDepthwiseConv2dNative"]:
             op_types = [op.decode() for op in node.get_attr_value("fused_ops")]
             extra_inputs = node.input[2:]
             g.replace_inputs(node, node.input[:2])
@@ -21,12 +21,14 @@ def rewrite_fused_ops(g, ops):
             shape = g.get_shape(node.output[0])
             first_node = None
             for op in op_types:
-                new_node = g.make_node(op, [last_output] + extra_inputs, skip_conversion=False,
+                num_inputs = {"BiasAdd": 2, "FusedBatchNorm": 5}.get(op, 1 + len(extra_inputs))
+                my_inputs = [last_output] + extra_inputs[:num_inputs - 1]
+                new_node = g.make_node(op, my_inputs, skip_conversion=False,
                                        op_name_scope=node.name, dtypes=[dtype], shapes=[shape])
                 last_output = new_node.output[0]
+                extra_inputs = extra_inputs[num_inputs - 1:]
                 if first_node is None:
                     first_node = new_node
-                extra_inputs = []
 
             consumers = [n for n in g.find_output_consumers(node.output[0]) if n != first_node]
             g.replace_all_inputs(node.output[0], last_output, consumers)

tf2onnx/tfjs_utils.py

@@ -100,15 +100,17 @@ def resolve_output(output, op_info, func_name=None):
         long_node_name = func_name + "/" + node
         if long_node_name in op_info:
             node = long_node_name
-    op_type, tf_attr = op_info[node]
-    names, _ = get_output_names_and_dtypes(op_type, tf_attr)
+    op_type, tf_attr, inp_dtypes = op_info[node]
+    names, _ = get_output_names_and_dtypes(op_type, tf_attr, inp_dtypes)
     idx = names.index(output_arg_name) + int(index)
     return node + ':' + str(idx)
 
 
-def get_output_names_and_dtypes(op_type, tf_attr):
+def get_output_names_and_dtypes(op_type, tf_attr, inp_dtypes):
     """Parses the tf documentation to determine the names and dtypes of the outputs of the op"""
     # TODO: ['Prelu', 'Conv1D', 'DepthwiseConv2d', 'FusedDepthwiseConv2dNative', 'Ones', 'Zeros']
+    if op_type == 'Prelu':
+        return ['activations'], [inp_dtypes[0]]
     try:
         tf_op_def = tf_api_def_map.get_op_def(op_type)
     except ValueError:
@@ -134,15 +136,19 @@ def get_output_names_and_dtypes(op_type, tf_attr):
     return names, dtypes
 
 
-def get_output_dtypes(op_type, tf_attr):
+def get_output_dtypes(op_type, tf_attr, inp_dtypes):
     """Returns a list of the tf dtypes for the op's outputs"""
-    _, out_dtypes = get_output_names_and_dtypes(op_type, tf_attr)
+    _, out_dtypes = get_output_names_and_dtypes(op_type, tf_attr, inp_dtypes)
     return out_dtypes
 
 
 def get_output_shapes(node_def, input_dtypes, input_shapes, inp_consts):
     """Returns a list of the output shapes of an op. input_dtypes should be tf dtypes."""
     from tf2onnx.tf_loader import tf_session, tf_placeholder  # pylint: disable=import-outside-toplevel
+
+    if node_def.op == "Prelu":
+        return [input_shapes[0]]
+
     del node_def.input[:]
     node_def.name = "node"
 
@@ -213,6 +219,16 @@ def graphs_from_tfjs(model_path, input_names=None, output_names=None, shape_over
     topologically sorted list of subgraphs."""
     model, zip_compressed = read_model_json(model_path)
 
+    model_format = model['modelTopology'].get('format')
+    if model_format is None:
+        if 'keras_version' in model['modelTopology']:
+            model_format = 'layers-model'
+        else:
+            model_format = 'graph-model'
+    utils.make_sure(model_format == 'graph-model', "tf2onnx only supports conversion from tfjs graph models, "
+                    "not format %s. Use Google's tfjs converter to convert to a graph model, then try again.",
+                    model_format)
+
     weights_manifest = model['weightsManifest'][0]
 
     sharded_data = []
@@ -356,7 +372,7 @@ def update_shapes(new_shapes):
             onnx_nodes.append(onnx_node)
             output_shapes[out_name] = shape_override.get(out_name, list(np_arr.shape))
             tf_dtypes[out_name] = tf_dtype
-            op_info[node_name] = (op_type, {'dtype': tf_dtypes[out_name]})
+            op_info[node_name] = (op_type, {'dtype': tf_dtypes[out_name]}, [])
             continue
         tf_attr = {}
         onnx_attr = {}
@@ -368,16 +384,25 @@ def update_shapes(new_shapes):
             if k == 'DstT':
                 k = 'to'
             onnx_attr[k] = read_tfjs_attr(v)
-        op_info[node_name] = (op_type, tf_attr)
+        if op_type == "FusedDepthwiseConv2dNative":
+            # This op isn't in tensorflow but can be converted to a TF op
+            op_type = "_FusedDepthwiseConv2dNative"
+            err_msg = "explicit_paddings for supported for _FusedDepthwiseConv2dNative"
+            utils.make_sure(len(tf_attr['explicit_paddings']) == 0, err_msg)
+            del tf_attr['explicit_paddings']
+            del onnx_attr['explicit_paddings']
+            del node_def.attr['explicit_paddings']
+            node_def.op = op_type
 
         input_names = [inp for inp in node.get('input', []) if not inp.startswith('^')]
         input_names = [resolve_output(inp, op_info, func_name) for inp in input_names]
         unused_outputs.difference_update(input_names)
         inp_dtypes = [tf_dtypes[inp] for inp in input_names]
         inp_shapes = [output_shapes[inp] for inp in input_names]
         inp_consts = [weights.get(inp.split(':')[0]) for inp in input_names]
-        out_dtypes = get_output_dtypes(op_type, tf_attr)
+        out_dtypes = get_output_dtypes(op_type, tf_attr, inp_dtypes)
         out_shapes = get_output_shapes(node_def, inp_dtypes, inp_shapes, inp_consts)
+        op_info[node_name] = (op_type, tf_attr, inp_dtypes)
 
         output_names = [node_name + ":" + str(i) for i in range(len(out_dtypes))]
         tf_dtypes.update(zip(output_names, out_dtypes))