[Experimental] Update accuracy checker for tf support and map input dims...

[Experimental] Update accuracy checker for tf support and map input dims (custom input shape) (#1514)

* update file reading function to fix external data loading

tools/accuracy/accuracy_checker.py

@@ -25,6 +25,7 @@ import argparse
 import numpy as np
 import migraphx
 import onnxruntime as ort
+import sys
 
 
 def parse_args():
@@ -33,15 +34,13 @@ def parse_args():
         'MIGraphX accuracy checker. Use to verify onnx files to ensure MIGraphX\'s output \
                                                   is within tolerance of onnx runtime\'s expected output.'
     )
-    req_args = parser.add_argument_group(title='required arguments')
-    req_args.add_argument('--onnx',
-                          type=str,
-                          required=True,
-                          help='path to onnx file')
-    req_args.add_argument('--provider',
-                          type=str,
-                          default='CPUExecutionProvider',
-                          help='execution provider for onnx runtime \
+    file_args = parser.add_argument_group(title='file type arguments')
+    file_args.add_argument('--onnx', type=str, help='path to onnx file')
+    file_args.add_argument('--tf', type=str, help='path to tf pb file')
+    parser.add_argument('--provider',
+                        type=str,
+                        default='CPUExecutionProvider',
+                        help='execution provider for onnx runtime \
                                 (default = CPUExecutionProvider)')
     parser.add_argument('--batch',
                         type=int,
@@ -50,6 +49,9 @@ def parse_args():
     parser.add_argument('--fill1',
                         action='store_true',
                         help='fill all arguments with a value of 1')
+    parser.add_argument('--fill0',
+                        action='store_true',
+                        help='fill all arguments with a value of 0')
     parser.add_argument('--verbose',
                         action='store_true',
                         help='show verbose information (for debugging)')
@@ -57,6 +59,12 @@ def parse_args():
                         type=float,
                         default=1e-3,
                         help='accuracy tolerance (default = 1e-3)')
+    parser.add_argument('--input-dim',
+                        type=str,
+                        action='append',
+                        help='specify input parameter dimension \
+                                with the following format --input_dim input_name:dim0,dim1,dim2...'
+                        )
     args = parser.parse_args()
 
     return args
@@ -111,42 +119,127 @@ def get_np_datatype(in_type):
 def main():
     args = parse_args()
 
+    use_onnx = True
+    if args.onnx == None:
+        use_onnx = False
+    if not use_onnx and args.tf == None:
+        print('Error: please specify either an onnx or tf pb file')
+        sys.exit(-1)
+
     model_name = args.onnx
+
     batch = args.batch
 
-    model = migraphx.parse_onnx(model_name, default_dim_value=batch)
+    custom_inputs = args.input_dim
+
+    input_dims = {}
+    if custom_inputs != None:
+        for input in custom_inputs:
+            input_dim = ''.join(input.split(':')[:-1])
+            dims = [int(dim) for dim in input.split(':')[-1].split(',')]
+            input_dims[input_dim] = dims
+
+    if use_onnx:
+        if not input_dims:
+            model = migraphx.parse_onnx(model_name, default_dim_value=batch)
+        else:
+            model = migraphx.parse_onnx(model_name,
+                                        default_dim_value=batch,
+                                        map_input_dims=input_dims)
+    else:
+        model_name = args.tf
+
+        if not input_dims:
+            model = migraphx.parse_tf(model_name, batch_size=batch)
+        else:
+            model = migraphx.parse_tf(model_name,
+                                      batch_size=batch,
+                                      map_input_dims=input_dims)
 
     if args.verbose:
         print(model)
 
-    model.compile(migraphx.get_target('gpu'), offload_copy=False)
+    model.compile(migraphx.get_target('gpu'))
 
     params = {}
     test_inputs = {}
     for name, shape in model.get_parameter_shapes().items():
         if args.verbose:
-            print('Parameter {} -> {}'.format(name, shape))
+            print(f'Parameter {name} -> {shape}')
         in_shape = shape.lens()
         in_type = shape.type_string()
-        if not args.fill1:
+        if not args.fill1 and not args.fill0:
             test_input = np.random.rand(*(in_shape)).astype(
                 get_np_datatype(in_type))
-        else:
+        elif not args.fill0:
             test_input = np.ones(in_shape).astype(get_np_datatype(in_type))
+        else:
+            test_input = np.zeros(in_shape).astype(get_np_datatype(in_type))
         test_inputs[name] = test_input
-        params[name] = migraphx.to_gpu(migraphx.argument(test_input))
+        params[name] = migraphx.argument(test_input)
+
+    pred_migx = np.array(model.run(params)[-1])
 
-    pred_migx = np.array(migraphx.from_gpu(model.run(params)[-1]))
+    if use_onnx:
+        sess = ort.InferenceSession(model_name, providers=[args.provider])
 
-    sess = ort.InferenceSession(model_name, providers=[args.provider])
+        ort_params = {}
+        for input in sess.get_inputs():
+            ort_params[input.name] = test_inputs[input.name]
+
+        try:
+            pred_fw = sess.run(None, ort_params)[-1]
+        except Exception as e:
+            if any(input_dims):
+                print(
+                    'Error: custom input dim may not be compatible with onnx runtime'
+                )
+            raise e
+    else:
+        import tensorflow as tf
+
+        def load_tf_graph(model_name):
+            with tf.io.gfile.GFile(model_name, 'rb') as f:
+                graph_def = tf.compat.v1.GraphDef()
+                graph_def.ParseFromString(f.read())
+
+            with tf.compat.v1.Graph().as_default() as graph:
+                tf.graph_util.import_graph_def(graph_def)
+            return graph
+
+        graph = load_tf_graph(model_name)
+        is_nhwc = False
+        graph_ops = []
+        for op in graph.get_operations():
+            graph_ops.append(op.name)
+            if 'Conv' in op.node_def.op:
+                if 'NHWC' in op.get_attr('data_format').decode('utf-8'):
+                    is_nhwc = True
+        graph_ops_set = set(graph_ops)
+        tf_dict = {}
+
+        for name in test_inputs.keys():
+            # graph.get_operations() adds 'import/' to the op name
+            tf_name = f'import/{name}'
+            if tf_name not in graph_ops_set:
+                continue
+            x = graph.get_tensor_by_name(f'{tf_name}:0')
+            tf_input = test_inputs[name]
+            # transpose input for NHWC model
+            if tf_input.ndim == 4 and is_nhwc:
+                tf_dict[x] = np.transpose(tf_input, (0, 2, 3, 1))
+            else:
+                tf_dict[x] = tf_input
 
-    ort_params = {}
-    for input in sess.get_inputs():
-        ort_params[input.name] = test_inputs[input.name]
+        # assume last node in graph is output
+        # TODO: let user specify op name for output
+        y = graph.get_tensor_by_name(f'{graph_ops[-1]}:0')
 
-    pred_ort = sess.run(None, ort_params)[-1]
+        with tf.compat.v1.Session(graph=graph) as sess:
+            y_out = sess.run(y, feed_dict=tf_dict)
+            pred_fw = y_out
 
-    is_correct = check_correctness(pred_ort, pred_migx, args.tolerance,
+    is_correct = check_correctness(pred_fw, pred_migx, args.tolerance,
                                    args.tolerance, args.verbose)
     verbose_string = ' Rerun with --verbose for detailed information.' \
             if not args.verbose else ''