Enhancements to qonnx-exec

src/qonnx/util/exec_qonnx.py

@@ -60,12 +60,18 @@ def exec_qonnx(
     expose_intermediates: str = None,
     output_prefix: str = "out_",
     output_mode: output_mode_options = OUTPUT_MODE_NAME,
-    argmax_verify_npy: str = None,
+    verify_npy: str = None,
+    verify_argmax=False,
     save_modified_model: str = None,
+    input_to_nchw=False,
+    input_to_nhwc=False,
+    input_cast2float=False,
     input_pix2float=False,
     input_zerocenter=False,
     maxiters: int = None,
-    output_nosave=False
+    output_nosave=False,
+    early_exit_acc_ratio=None,
+    override_exec_onnx=None
 ):
     """Execute a given QONNX model by initializing its inputs from .npy files, and write outputs
     as .npy files.
@@ -79,13 +85,19 @@ def exec_qonnx(
         Matched patterns will expose intermediate outputs as top-level outputs.
     :param output_prefix: Prefix for the generated output files.
     :param output_mode: Naming mode for generated output files.
-    :param argmax_verify_npy: If specified, take argmax of output and compare to this file for top-1 accuracy measurement
+    :param verify_npy: If specified, compare output to this file for top-1 accuracy measurement
+    :param verify_argmax: If specified, take argmax of output before comparing to verify_npy
     :param save_modified_model: If specified, save the modified model
         (after batchsize changes or exposed intermediate tensors) with this filename
+    :param input_to_nchw: If specified, convert input tensors to NCHW format
+    :param input_to_nhwc: If specified, convert input tensors to NHWC format
+    :param input_cast2float: If specified, apply simple cast to float32 for input
     :param input_pix2float: If specified, do uint8 [0,255] -> fp32 [0,1] mapping for input
     :param input_zerocenter: If specified together with pix2float, do uint8 [0,255] -> fp32 [-1,+1] mapping for input
     :param maxiters: If specified, limit maximum number of iterations (batches) to be processed
     :param output_nosave: If specified, do not save output tensors to files
+    :param early_exit_acc_ratio: If specified as a float number between 0 and 1, early exit if any batch accuracy falls under
+    :param override_exec_onnx: If specified, use this function to execute the model instead of qonnx
     """
     assert output_mode in output_modes, "Unrecognized output mode"
 
@@ -132,8 +144,8 @@ def exec_qonnx(
             inp = inp.reshape(n_dset_iters, bsize, *inp.shape[1:])
             inp_data_reshaped.append(inp)
         inp_data = inp_data_reshaped
-        if argmax_verify_npy is not None:
-            labels = np.load(argmax_verify_npy)
+        if verify_npy is not None:
+            labels = np.load(verify_npy)
             assert labels.shape[0] == dset_size, "Label size must match dataset size"
             labels = labels.reshape(n_dset_iters, bsize, *labels.shape[1:])
     else:
@@ -150,24 +162,34 @@ def exec_qonnx(
     for iter in pbar:
         iter_suffix = "_batch%d" % iter
         idict = {}
-        if not argmax_verify_npy:
+        if not verify_npy:
             pbar.set_description("Batch [%d/%d]: running" % (iter + 1, n_dset_iters))
         # supply inputs and execute
         for inp_ind, inp in enumerate(model.graph.input):
             if input_pix2float:
                 idict[inp.name] = (inp_data[inp_ind][iter] / 255.0).astype(np.float32)
                 if input_zerocenter:
                     idict[inp.name] = (2 * idict[inp.name] - 1.0).astype(np.float32)
+            elif input_cast2float:
+                idict[inp.name] = inp_data[inp_ind][iter].astype(np.float32)
             else:
                 idict[inp.name] = inp_data[inp_ind][iter]
-        if n_custom_nodes > 0:
-            # run node-by-node in qonnx
-            odict = execute_onnx(model, idict)
+            if input_to_nhwc:
+                idict[inp.name] = np.transpose(idict[inp.name], (0, 2, 3, 1))
+            if input_to_nchw:
+                idict[inp.name] = np.transpose(idict[inp.name], (0, 3, 1, 2))
+        if override_exec_onnx is not None:
+            # run using specified custom execution function
+            odict = override_exec_onnx(model, idict)
         else:
-            # run using onnxruntime
-            sess = rt.InferenceSession(model.model.SerializeToString())
-            output_list = sess.run(None, idict)
-            odict = {outp.name: output_list[oind] for oind, outp in enumerate(model.graph.output)}
+            if n_custom_nodes > 0:
+                # run node-by-node in qonnx
+                odict = execute_onnx(model, idict)
+            else:
+                # run using onnxruntime
+                sess = rt.InferenceSession(model.model.SerializeToString())
+                output_list = sess.run(None, idict)
+                odict = {outp.name: output_list[oind] for oind, outp in enumerate(model.graph.output)}
         if not output_nosave:
             for out_ind, outp in enumerate(model.graph.output):
                 # save generated outputs
@@ -176,10 +198,11 @@ def exec_qonnx(
                 elif output_mode == OUTPUT_MODE_NAME:
                     oname = outp.name
                 np.save(output_prefix + oname + iter_suffix + ".npy", odict[outp.name])
-        if argmax_verify_npy:
+        if verify_npy:
             # measure accuracy for output
             ret = odict[model.graph.output[0].name]
-            ret = np.argmax(ret, axis=-1)
+            if verify_argmax:
+                ret = np.argmax(ret, axis=-1)
             ok_batch = np.count_nonzero(ret == labels[iter])
             nok_batch = bsize - ok_batch
             ok += ok_batch
@@ -190,6 +213,9 @@ def exec_qonnx(
                 "Batch [%d/%d]: ok %d nok %d accuracy %f (overall ok %d nok %d accuracy %f)"
                 % (iter + 1, n_dset_iters, ok_batch, nok_batch, accuracy_batch, ok, nok, accuracy_overall)
             )
+            if early_exit_acc_ratio is not None and accuracy_batch < early_exit_acc_ratio:
+                return (ok, nok)
+    return (ok, nok)
 
 
 def main():